diff --git a/AGENTS.md b/AGENTS.md index 1011436d2..98212c48f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -535,7 +535,7 @@ Core retrieval internals are grouped by ownership: - `execution/`: request shaping, route selection (classic / agent_explore / small_corpus), and public response projection. - `search/`: `map_unit_discovery` (persisted map-unit BM25 discovery; incomplete or incompatible indexes raise), scoring, section filters, candidate ranking. - `hydration/`: row/path/reference hydration, inline assets, and result assembly. -- `agent_explore/` + `agent_tools/`: agent retrieval when `use_agentic` is on. Harness is `cursor_sdk` (Cursor) or `openai` (Knowhere's own). Map-nav is archived under `deprecated/mapnav/`. +- `agent_explore/` + `agent_tools/`: agent retrieval when `use_agentic` is on. Harness is `cursor_sdk` (Cursor) or `openai` (Knowhere's own). - `trace/`: `DecisionTraceStep` mapping and `TraceRecorder`. - `graph/`: document graph publication/query support. - `stats/`: retrieval hit recording. @@ -571,8 +571,7 @@ There is no scored term channel in this primary path. Chunk Default agentic path is `agent_explore`: a `corpus.*` tool loop run by the resolved harness (`AGENT_EXPLORE_HARNESS`, default `cursor_sdk`). Finish refs -are resolved to chunks, then assembled like classic. Map-nav is archived at -`deprecated/mapnav/` and is no longer a live route. +are resolved to chunks, then assembled like classic. ### Result Assembly diff --git a/CONTEXT.md b/CONTEXT.md index df1f3799e..1767b5a72 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -371,7 +371,7 @@ scope, filters, data type, ranking options, and agentic toggle. ### Retrieval Run -One execution of a Retrieval Query, recorded across classic, map-nav, +One execution of a Retrieval Query, recorded across classic, agent_explore, small-corpus, cache-hit, failed, and cancelled outcomes with route, timing, and terminal-status metadata. diff --git a/apps/api/.env.example b/apps/api/.env.example index 06114302e..ddd2c4ebc 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -85,6 +85,7 @@ GPT_API_KEY= GLM_API_KEY= ALI_API_KEYS= ARK_API_KEY= +VERCEL_KEY= # Optional overrides: # DS_URL=https://api.deepseek.com/v1 @@ -97,6 +98,9 @@ ARK_API_KEY= # ASSET_MODEL is optional; empty follows IMAGE_MODEL. # IMAGE_MODEL=deepseek-flash # ASSET_MODEL= +# Page-image and query embeddings use VERCEL_KEY. +# EMBEDDING_MODEL=google/gemini-embedding-2 +# EMBEDDING_DIM=3072 # Optional retrieval overrides have code defaults. Retrieval is evidence-only: # evidence is the primary output; evidence_text is deprecated and always empty; diff --git a/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py b/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py index 300dc99ca..2614e6303 100644 --- a/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py +++ b/apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py @@ -1,4 +1,4 @@ -"""Add the index used by lazy map-nav section loads.""" +"""Add the index used by revision-section chunk loads.""" from __future__ import annotations diff --git a/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py index 79aa7583c..03e4c49f0 100644 --- a/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py +++ b/apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py @@ -1,4 +1,4 @@ -"""Add revision-pinned map-nav score units.""" +"""Add revision-pinned map-unit score rows.""" from __future__ import annotations diff --git a/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py index 61250639a..b0c2cce1c 100644 --- a/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py +++ b/apps/api/alembic/versions/2e3f4a5b6c7d_add_map_unit_token_covering_index.py @@ -1,4 +1,4 @@ -"""Add a covering index for persisted map-nav token lookups.""" +"""Add a covering index for persisted map-unit token lookups.""" from __future__ import annotations diff --git a/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py index 4dcc50117..c0a120cbe 100644 --- a/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py +++ b/apps/api/alembic/versions/3f4a5b6c7d8e_add_section_snapshot_order_index.py @@ -1,4 +1,4 @@ -"""Add the index used by lazy map-nav section pagination.""" +"""Add the index used by section snapshot-order pagination.""" from __future__ import annotations diff --git a/apps/api/alembic/versions/6f7a8b9c0d1e_add_document_scope_enrichments.py b/apps/api/alembic/versions/6f7a8b9c0d1e_add_document_scope_enrichments.py new file mode 100644 index 000000000..8797acd0f --- /dev/null +++ b/apps/api/alembic/versions/6f7a8b9c0d1e_add_document_scope_enrichments.py @@ -0,0 +1,51 @@ +"""add document scope enrichments + +Revision ID: 6f7a8b9c0d1e +Revises: 5e6f7a8b9c0d +Create Date: 2026-10-26 09:30:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "6f7a8b9c0d1e" +down_revision: str = "5e6f7a8b9c0d" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "document_scope_enrichments", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("scope_id", sa.String(length=64), nullable=False), + sa.Column("page_start", sa.Integer(), nullable=False), + sa.Column("page_end", sa.Integer(), nullable=False), + sa.Column("state", sa.String(length=32), nullable=False), + sa.Column("priority", sa.String(length=32), nullable=True), + sa.Column("attempts", sa.Integer(), nullable=False), + sa.Column("last_hit_at", sa.DateTime(), nullable=True), + sa.Column("refined_at", sa.DateTime(), nullable=True), + sa.Column("updated_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "job_result_id", + "scope_id", + name="uq_document_scope_enrichments_revision_scope", + ), + ) + + +def downgrade() -> None: + op.drop_table("document_scope_enrichments") diff --git a/apps/api/alembic/versions/7a8b9c0d1e2f_add_document_page_embeddings.py b/apps/api/alembic/versions/7a8b9c0d1e2f_add_document_page_embeddings.py new file mode 100644 index 000000000..4640b18e4 --- /dev/null +++ b/apps/api/alembic/versions/7a8b9c0d1e2f_add_document_page_embeddings.py @@ -0,0 +1,136 @@ +"""add document page embeddings + +Revision ID: 7a8b9c0d1e2f +Revises: 6f7a8b9c0d1e +Create Date: 2026-10-26 17:10:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector + +revision: str = "7a8b9c0d1e2f" +down_revision: str = "6f7a8b9c0d1e" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + op.execute("CREATE EXTENSION IF NOT EXISTS vector") + + op.create_table( + "document_page_embeddings", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("chunk_id", sa.String(length=64), nullable=False), + sa.Column("page_num", sa.Integer(), nullable=False), + sa.Column("modality", sa.String(length=16), nullable=False), + sa.Column("model", sa.String(length=64), nullable=False), + sa.Column("dims", sa.Integer(), nullable=False), + sa.Column("embedding", Vector(1024), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "document_id", + "job_result_id", + "page_num", + name="uq_document_page_embeddings_revision_page", + ), + ) + op.create_index( + "idx_document_page_embeddings_scope", + "document_page_embeddings", + ["user_id", "namespace", "job_result_id"], + ) + op.create_index( + "idx_document_page_embeddings_revision", + "document_page_embeddings", + ["document_id", "job_result_id", "page_num"], + ) + op.create_index( + "idx_document_page_embeddings_hnsw_cosine", + "document_page_embeddings", + ["embedding"], + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + + op.create_table( + "demo_document_page_embeddings", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("user_id", sa.Text(), nullable=False), + sa.Column("namespace", sa.String(length=255), nullable=False), + sa.Column("document_id", sa.String(length=36), nullable=False), + sa.Column("job_result_id", sa.String(length=36), nullable=False), + sa.Column("chunk_id", sa.String(length=64), nullable=False), + sa.Column("page_num", sa.Integer(), nullable=False), + sa.Column("modality", sa.String(length=16), nullable=False), + sa.Column("model", sa.String(length=64), nullable=False), + sa.Column("dims", sa.Integer(), nullable=False), + sa.Column("embedding", Vector(1024), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], ["demo_documents.document_id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["job_result_id"], ["job_results.id"], ondelete="RESTRICT" + ), + sa.ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_documentpageembedding_revision", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "document_id", + "job_result_id", + "page_num", + name="uq_demo_document_page_embeddings_revision_page", + ), + ) + op.create_index( + "idx_demo_document_page_embeddings_scope", + "demo_document_page_embeddings", + ["user_id", "namespace", "job_result_id"], + ) + op.create_index( + "idx_demo_document_page_embeddings_revision", + "demo_document_page_embeddings", + ["document_id", "job_result_id", "page_num"], + ) + op.create_index( + "idx_demo_document_page_embeddings_hnsw_cosine", + "demo_document_page_embeddings", + ["embedding"], + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + op.execute("ALTER TABLE demo_document_page_embeddings ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE demo_document_page_embeddings FORCE ROW LEVEL SECURITY") + op.execute( + "CREATE POLICY demo_maintainer_write ON demo_document_page_embeddings " + "FOR ALL USING (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '') " + "WITH CHECK (COALESCE(current_setting('knowhere.demo_maintainer', true), '') <> '')" + ) + op.execute( + "CREATE POLICY demo_public_read ON demo_document_page_embeddings " + "FOR SELECT USING (is_demo_revision_ready(document_id, job_result_id))" + ) + + +def downgrade() -> None: + op.drop_table("demo_document_page_embeddings") + op.drop_table("document_page_embeddings") diff --git a/apps/api/alembic/versions/8b9c0d1e2f3a_resize_page_embeddings_to_3072.py b/apps/api/alembic/versions/8b9c0d1e2f3a_resize_page_embeddings_to_3072.py new file mode 100644 index 000000000..7ec0cf1ac --- /dev/null +++ b/apps/api/alembic/versions/8b9c0d1e2f3a_resize_page_embeddings_to_3072.py @@ -0,0 +1,60 @@ +"""resize page embeddings to 3072 + +Revision ID: 8b9c0d1e2f3a +Revises: 7a8b9c0d1e2f +Create Date: 2026-10-26 21:40:00.000000 + +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "8b9c0d1e2f3a" +down_revision: str = "7a8b9c0d1e2f" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +_TABLES = ( + "document_page_embeddings", + "demo_document_page_embeddings", +) + + +def _resize_embedding_column( + table_name: str, + *, + column_type: str, + index_ops: str, +) -> None: + index_name = f"idx_{table_name}_hnsw_cosine" + op.drop_index(index_name, table_name=table_name) + op.execute(f"DELETE FROM {table_name}") + op.execute( + f"ALTER TABLE {table_name} ALTER COLUMN embedding TYPE {column_type}" + ) + op.create_index( + index_name, + table_name, + ["embedding"], + postgresql_using="hnsw", + postgresql_ops={"embedding": index_ops}, + ) + + +def upgrade() -> None: + for table_name in _TABLES: + _resize_embedding_column( + table_name, + column_type="halfvec(3072)", + index_ops="halfvec_cosine_ops", + ) + + +def downgrade() -> None: + for table_name in reversed(_TABLES): + _resize_embedding_column( + table_name, + column_type="vector(1024)", + index_ops="vector_cosine_ops", + ) diff --git a/apps/api/app/services/document_ingestion/command.py b/apps/api/app/services/document_ingestion/command.py index cb905eae5..f281ad2a4 100644 --- a/apps/api/app/services/document_ingestion/command.py +++ b/apps/api/app/services/document_ingestion/command.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from typing import Literal -from shared.models.schemas.job import JobCreateBase +from shared.models.schemas.job import IngestionMode, JobCreateBase from shared.models.schemas.page_memory_config import PageMemoryConfig ApiVersion = Literal["v1", "v2"] @@ -20,6 +20,7 @@ class DocumentIngestionCommand: api_version: ApiVersion parse_track: ParseTrack processing_generation: ProcessingGeneration + ingestion_mode: IngestionMode = "full" page_memory_config: PageMemoryConfig | None = None @@ -57,11 +58,13 @@ def _build_ingestion_command( page_memory_config = ( PageMemoryConfig.default() if parse_track == "page_memory" else None ) + ingestion_mode: IngestionMode = getattr(payload, "ingestion_mode", "full") return DocumentIngestionCommand( payload=payload, api_version=api_version, parse_track=parse_track, processing_generation=processing_generation, + ingestion_mode=ingestion_mode, page_memory_config=page_memory_config, ) diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index fc2a0e6dc..77a1f50d7 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -122,6 +122,19 @@ async def _create_job_with_command_factory( api_version=api_version, ) command = build_command(file_extension) + if command.ingestion_mode == "progressive" and command.parse_track != "page_memory": + raise ValidationException( + user_message="Progressive ingestion is only supported for page-memory formats (.pdf, .pptx) on /v2/jobs", + violations=[ + { + "field": "ingestion_mode", + "description": ( + "ingestion_mode=progressive requires a page_memory " + "document (.pdf, .pptx) submitted via /v2/jobs" + ), + } + ], + ) scope = await self._resolve_scope( db, command=command, @@ -284,6 +297,7 @@ async def _resolve_scope( parse_track=command.parse_track, processing_generation=command.processing_generation, page_memory_config=command.page_memory_config, + ingestion_mode=command.ingestion_mode, ), ) if CorpusStorage.resolve_namespace(payload.namespace).is_demo: diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index 48d9bef16..2688547d6 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -1,4 +1,4 @@ -"""Backfill persisted MAP-NAV lexical indexes for existing revisions. +"""Backfill persisted map-unit lexical indexes for existing revisions. Rebuilds, per active revision: the map-unit index, the revision serving manifest, that document's subtree in the namespace MAP snapshot, and the @@ -88,7 +88,7 @@ def _bootstrap_python_path() -> None: def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( - description="Backfill MAP-NAV indexes for current document revisions." + description="Backfill map-unit indexes for current document revisions." ) parser.add_argument( "--apply", diff --git a/apps/api/scripts/ensure_test_environment.py b/apps/api/scripts/ensure_test_environment.py index b52719aa8..47f1cffa3 100644 --- a/apps/api/scripts/ensure_test_environment.py +++ b/apps/api/scripts/ensure_test_environment.py @@ -43,15 +43,18 @@ class InstallCommand: _REQUIRED_POSTGRESQL_EXTENSIONS: tuple[str, ...] = ( "uuid-ossp", "pg_trgm", + "vector", ) def _get_postgresql_install_hint() -> str: return ( "Run `uv run python apps/api/scripts/ensure_test_environment.py --install`, " - "or install PostgreSQL server binaries and contrib extensions manually. " + "or install PostgreSQL server binaries, contrib extensions, and pgvector. " "Ubuntu: sudo apt-get update && sudo apt-get install -y " - "postgresql postgresql-contrib. macOS: brew install postgresql@16 " + "postgresql postgresql-contrib, then install " + "postgresql--pgvector from PGDG without upgrading " + "the PostgreSQL metapackage. macOS: brew install postgresql@16 pgvector " "and add its bin directory to PATH. If needed, set " "PYTEST_POSTGRESQL_EXECUTABLE to the pg_ctl path." ) @@ -205,13 +208,31 @@ def get_install_commands() -> tuple[InstallCommand, ...]: ("apt-get", "install", "-y", "postgresql", "postgresql-contrib") ), ), + InstallCommand( + description=( + "Hold Ubuntu PostgreSQL and install matching PGDG pgvector" + ), + command=_with_privilege( + ( + "bash", + "-lc", + ( + "set -euo pipefail; " + "major=$(pg_config --version | awk '{print $2}' | cut -d. -f1); " + "/usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y; " + 'apt-mark hold postgresql postgresql-contrib "postgresql-${major}"; ' + 'apt-get install -y --no-upgrade "postgresql-${major}-pgvector"' + ), + ) + ), + ), ) if system_name == "darwin" and shutil.which("brew") is not None: return ( InstallCommand( description="Install PostgreSQL server binaries and contrib extensions", - command=("brew", "install", "postgresql@16"), + command=("brew", "install", "postgresql@16", "pgvector"), ), ) diff --git a/apps/api/tests/README.md b/apps/api/tests/README.md index cdd7d84a1..b55be7469 100644 --- a/apps/api/tests/README.md +++ b/apps/api/tests/README.md @@ -27,7 +27,7 @@ This directory is the new home for API tests. ## Local Prerequisites - PostgreSQL server binaries available to `pytest-postgresql` -- PostgreSQL contrib extension files for `uuid-ossp` and `pg_trgm` +- PostgreSQL contrib extension files for `uuid-ossp`, `pg_trgm`, and `vector` - No running local PostgreSQL service is required; `pytest-postgresql` starts an isolated process - No running local Redis service is required; contract tests use `fakeredis` - Set `PYTEST_POSTGRESQL_EXECUTABLE` to the `pg_ctl` path if discovery cannot find it diff --git a/apps/api/tests/contract/test_page_memory_parse_track_contract.py b/apps/api/tests/contract/test_page_memory_parse_track_contract.py index 17e7752a1..09c92e015 100644 --- a/apps/api/tests/contract/test_page_memory_parse_track_contract.py +++ b/apps/api/tests/contract/test_page_memory_parse_track_contract.py @@ -171,10 +171,45 @@ def test_public_job_create_schemas_do_not_advertise_parse_track() -> None: assert "parse_track" not in v1_schema["properties"] assert "parse_track" not in v2_schema["properties"] + assert "ingestion_mode" not in v1_schema["properties"] + assert "ingestion_mode" in v2_schema["properties"] assert v1_schema["additionalProperties"] is False assert v2_schema["additionalProperties"] is False +def test_v2_progressive_ingestion_mode_rejected_for_non_page_memory_formats() -> None: + import asyncio + from unittest.mock import AsyncMock + + from app.services.document_ingestion.service import DocumentIngestionService + from app.services.rate_limit.data_structures import CurrentUser + from shared.models.schemas.job_metadata import JobMetadataHelper + + payload = _build_job_create_v2( + source_type="file", + file_name="policy.docx", + ingestion_mode="progressive", + ) + meta = JobMetadataHelper.create_from_request(payload, api_version="v2") + assert JobMetadataHelper.get_ingestion_mode(meta) == "progressive" + assert JobMetadataHelper.get_ingestion_mode({}) == "full" + + svc = DocumentIngestionService( + creation_service=AsyncMock(), + confirmation_service=AsyncMock(), + job_admission_service=AsyncMock(), + ) + with pytest.raises(ValidationException) as exc_info: + asyncio.run( + svc.create_v2_job( + AsyncMock(), + payload=payload, + current_user=CurrentUser(user_id="u1", user_tier="pro"), + ) + ) + assert exc_info.value.details["violations"][0]["field"] == "ingestion_mode" + + def test_v2_job_polling_system_rule_matches_before_default() -> None: from app.services.rate_limit.data_structures import SystemLimitRule from app.services.rate_limit.system_limit import find_system_rule diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 8eb5e4b05..69a9f9236 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -508,46 +508,6 @@ async def test_agent_explore_retrieval_should_return_seeded_chunk_via_fake_episo } -@pytest.mark.asyncio -async def test_agentic_router_env_mapnav_is_ignored( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: MonkeyPatch, -) -> None: - monkeypatch.setenv("RETRIEVAL_AGENTIC_ROUTER", "mapnav") - async with developer_api_client_factory() as api_client: - target = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-explore-env", - source_file_name="target.pdf", - section_path="Findings", - content="env ignored content", - ) - await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-explore-env", - source_file_name="filler.pdf", - section_path="filler/section", - content="unrelated filler content", - ) - _patch_harness( - monkeypatch, - _episode_keeping_refs(documents=[target], notes="env ignored"), - ) - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-explore-env", - "query": "ignored", - "top_k": 1, - "use_agentic": True, - }, - ) - assert response.status_code == 200 - assert cast(dict[str, object], response.json())["router_used"] == "agent_explore" - - @pytest.mark.asyncio async def test_agent_explore_should_not_hydrate_references_outside_request_scope( developer_api_client_factory: Callable[ diff --git a/apps/api/tests/contract/test_scope_enrichment_publish_contract.py b/apps/api/tests/contract/test_scope_enrichment_publish_contract.py new file mode 100644 index 000000000..d1c151bde --- /dev/null +++ b/apps/api/tests/contract/test_scope_enrichment_publish_contract.py @@ -0,0 +1,1550 @@ +"""Progressive publication: scope rows, queueing, read promotion, and any-stage tools.""" + +from __future__ import annotations + +import asyncio + +from sqlalchemy import create_engine, select +from sqlalchemy.engine import make_url +from sqlalchemy.orm import Session + +from shared.core.state_machine.transition_outcome import JobTransitionOutcome +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentMapUnitIndex, + DocumentScopeEnrichment, + DocumentSection, +) +from shared.services.jobs.lifecycle.scope_enrichments import ( + BUILD_STRUCTURE_TASK_NAME, + ENRICH_SCOPE_TASK_NAME, + HIGH_ENRICHMENT_QUEUE, + HIGH_ENRICHMENT_TASK_PRIORITY, + NORMAL_ENRICHMENT_QUEUE, + write_coarse_scope_enrichments, +) +from shared.services.jobs.lifecycle.success_finalizer import ( + JobSuccessFinalization, + SyncJobSuccessFinalizer, +) +from shared.services.retrieval.agent_tools.registry import ToolContext +from shared.services.retrieval.agent_tools.tools.assets import assets +from shared.services.retrieval.agent_tools.tools.grep import grep +from shared.services.retrieval.agent_tools.tools.list_documents import list_documents +from shared.services.retrieval.agent_tools.tools.node_filter import node_filter +from shared.services.retrieval.agent_tools.tools.outline import outline +from shared.services.retrieval.agent_tools.tools.query_table import query_table +from shared.services.retrieval.agent_tools.tools.read import read +from shared.services.retrieval.agent_tools.tools.recall import recall +from shared.services.retrieval.hydration.result_assembly import ( + assemble_retrieval_results, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery +from shared.testing.contract_runtime import ( + PostgreSQLProcess, + configure_contract_environment, + get_contract_database_url, + prepare_contract_storage, +) +from tests.support.publication_benchmark_support import ensure_benchmark_import_path + +ensure_benchmark_import_path() + +from scripts.publication_benchmark.publication_execution import ( # noqa: E402 + PublicationScope, + seed_publication_fixture, +) + + +class _DoneStateMachine: + def mark_completed_outcome( + self, + db: object, + job_id: str, + result_metadata: object = None, + operator_id: object = None, + ) -> JobTransitionOutcome: + del db, result_metadata, operator_id + return JobTransitionOutcome.transitioned( + job_id=job_id, + from_state="running", + to_state="done", + attempts=1, + ) + + +def _scope(label: str) -> PublicationScope: + return PublicationScope( + user_ref=f"scope-enrich-{label}-user", + namespace_ref=f"scope-enrich-{label}", + source_file_name=f"scope-enrich-{label}.pdf", + job_ref=f"scope-enrich-{label}-job", + revision_ref=f"scope-enrich-{label}-result", + ) + + +def _page_chunk( + *, + chunk_id: str, + path: str, + scope_id: str, + page_nums: list[int], +) -> dict[str, object]: + return { + "chunk_id": chunk_id, + "type": "page", + "content": f"raw {scope_id}", + "path": path, + "metadata": { + "scope_id": scope_id, + "enrichment_state": "coarse", + "page_nums": page_nums, + }, + } + + +def _set_job_progressive( + session: Session, + *, + job_id: str, + ingestion_mode: str = "progressive", + enrichment_status: str = "coarse", +) -> None: + from shared.models.database.job import Job + + job = session.execute(select(Job).where(Job.job_id == job_id)).scalar_one() + job.job_metadata = { + **dict(job.job_metadata or {}), + "parse_track": "page_memory", + "ingestion_mode": ingestion_mode, + "enrichment_status": enrichment_status, + } + session.flush() + + +def _finalize( + session: Session, + *, + job_id: str, + chunks: list[dict[str, object]], +) -> JobSuccessFinalization: + return SyncJobSuccessFinalizer(state_machine=_DoneStateMachine()).finalize( + session, + job_id=job_id, + result_s3_key="results/scope-enrich/result.zip", + checksum="checksum", + zip_size=1, + chunks=chunks, + stored_count=len(chunks), + delivery_mode="url", + section_summaries=None, + ) + + +def _write_scope_rows(session: Session, *, scope: PublicationScope) -> None: + document_id = session.execute( + select(Document.document_id).where( + Document.current_job_result_id == scope.revision_ref + ) + ).scalar_one() + assert write_coarse_scope_enrichments( + session, + document_id=document_id, + job_result_id=scope.revision_ref, + ) + + +def test_write_coarse_scope_enrichments_inserts_one_row_per_published_scope( + monkeypatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("page") + chunks = [ + _page_chunk( + chunk_id="coarse-a1", + path=f"{scope.source_file_name}/A", + scope_id="p1-2", + page_nums=[1, 2], + ), + _page_chunk( + chunk_id="coarse-a2", + path=f"{scope.source_file_name}/A2", + scope_id="p1-2", + page_nums=[1, 2], + ), + _page_chunk( + chunk_id="coarse-b", + path=f"{scope.source_file_name}/B", + scope_id="p3-4", + page_nums=[3, 4], + ), + ] + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + _set_job_progressive(session, job_id=scope.job_ref) + session.commit() + _finalize(session, job_id=scope.job_ref, chunks=chunks) + _write_scope_rows(session, scope=scope) + session.commit() + + with Session(engine) as session: + rows = ( + session.execute( + select(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == scope.revision_ref) + .order_by(DocumentScopeEnrichment.scope_id) + ) + .scalars() + .all() + ) + assert [ + (row.scope_id, row.page_start, row.page_end, row.state) for row in rows + ] == [ + ("p1-2", 1, 2, "coarse"), + ("p3-4", 3, 4, "coarse"), + ] + assert {row.document_id for row in rows} == {rows[0].document_id} + assert all(row.priority is None for row in rows) + assert all(row.attempts == 0 for row in rows) + assert all(row.refined_at is None for row in rows) + + with Session(engine) as session: + _write_scope_rows(session, scope=scope) + session.commit() + + with Session(engine) as session: + replay_count = len( + session.execute( + select(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == scope.revision_ref + ) + ) + .scalars() + .all() + ) + assert replay_count == 2 + finally: + engine.dispose() + + +def test_finalize_job_success_skips_enrichments_without_scope_id( + monkeypatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("text") + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + session.commit() + _finalize( + session, + job_id=scope.job_ref, + chunks=[ + { + "chunk_id": "text-body", + "type": "text", + "content": "plain text", + "path": f"{scope.source_file_name}/Root/Body", + "metadata": {}, + } + ], + ) + session.commit() + + with Session(engine) as session: + rows = ( + session.execute( + select(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == scope.revision_ref + ) + ) + .scalars() + .all() + ) + assert rows == [] + finally: + engine.dispose() + + +def test_result_assembly_branches_explicitly_on_page_enrichment_state() -> None: + rows = [ + { + "chunk_id": "page-text", + "document_id": "doc-1", + "chunk_type": "page", + "content": "raw m0 page text", + "chunk_metadata": { + "scope_id": "p1-1", + "enrichment_state": "text", + "summary": "", + "page_nums": [1], + }, + }, + { + "chunk_id": "page-coarse", + "document_id": "doc-1", + "chunk_type": "page", + "content": "raw page one text", + "chunk_metadata": { + "scope_id": "p1-2", + "enrichment_state": "coarse", + "summary": "", + "page_nums": [1, 2], + }, + }, + { + "chunk_id": "page-queued", + "document_id": "doc-1", + "chunk_type": "page", + "content": "raw page three text", + "chunk_metadata": { + "scope_id": "p3-4", + "enrichment_state": "queued", + "summary": "ignored while queued", + "page_nums": [3, 4], + }, + }, + { + "chunk_id": "page-refined-with-summary", + "document_id": "doc-1", + "chunk_type": "page", + "content": "raw page five text", + "chunk_metadata": { + "scope_id": "p5-6", + "enrichment_state": "refined", + "summary": "refined vlm summary", + "page_nums": [5, 6], + }, + }, + { + "chunk_id": "page-refined-empty-summary", + "document_id": "doc-1", + "chunk_type": "page", + "content": "raw text must not leak on refined", + "chunk_metadata": { + "scope_id": "p7-8", + "enrichment_state": "refined", + "summary": "", + "page_nums": [7, 8], + }, + }, + ] + + assembled = asyncio.run( + assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + ) + ) + by_id = {row["chunk_id"]: row for row in assembled} + + assert by_id["page-text"]["content"] == "raw m0 page text" + assert by_id["page-text"]["content_source"] == "raw" + assert by_id["page-text"]["page_nums"] == [1] + assert by_id["page-text"]["composed"] == [ + {"type": "text", "text": "raw m0 page text"} + ] + + assert by_id["page-coarse"]["content"] == "raw page one text" + assert by_id["page-coarse"]["content_source"] == "raw" + assert by_id["page-coarse"]["page_nums"] == [1, 2] + assert by_id["page-coarse"]["composed"] == [ + {"type": "text", "text": "raw page one text"} + ] + + assert by_id["page-queued"]["content"] == "raw page three text" + assert by_id["page-queued"]["content_source"] == "raw" + + assert by_id["page-refined-with-summary"]["content"] == "refined vlm summary" + assert by_id["page-refined-with-summary"]["content_source"] == "summary" + + # Explicit branch: refined never falls back to raw content when summary is empty. + assert by_id["page-refined-empty-summary"]["content"] == "" + assert by_id["page-refined-empty-summary"]["content_source"] == "summary" + + +def test_replace_document_scope_content_isolates_scope_reuses_ancestor_rebuilds_map_units_and_bumps_cache( + monkeypatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + from shared.core.database import get_db_context + from shared.services.redis import redis_sync_service + from shared.services.retrieval.cache_service import ( + get_retrieval_namespace_cache_version as get_cache_version, + ) + from shared.services.retrieval.publication_scope_content import ( + replace_document_scope_content as replace_scope_content, + ) + from shared.testing.contract_runtime import ( + CONTRACT_REDIS_DATABASE, + CONTRACT_REDIS_HOST, + CONTRACT_REDIS_PORT, + ContractFakeSyncRedis, + ) + + monkeypatch.setattr( + redis_sync_service, + "SyncRedisClient", + lambda **_kw: ContractFakeSyncRedis.from_url( + f"redis://{CONTRACT_REDIS_HOST}:{CONTRACT_REDIS_PORT}/{CONTRACT_REDIS_DATABASE}" + ), + ) + redis_sync_service.SyncRedisServiceFactory.reset() + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("partial") + coarse_chunks = [ + _page_chunk( + chunk_id="coarse-p1-2", + path=f"{scope.source_file_name}/SharedParent/OldLeafA", + scope_id="p1-2", + page_nums=[1, 2], + ) + | {"content": "alpha coarse unique token"}, + _page_chunk( + chunk_id="coarse-p3-4", + path=f"{scope.source_file_name}/SharedParent/SiblingLeafB", + scope_id="p3-4", + page_nums=[3, 4], + ) + | {"content": "bravo sibling untouched token"}, + _page_chunk( + chunk_id="coarse-p5-6", + path=f"{scope.source_file_name}/OtherChapter/LeafC", + scope_id="p5-6", + page_nums=[5, 6], + ) + | {"content": "charlie third scope token"}, + ] + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + _set_job_progressive(session, job_id=scope.job_ref) + session.commit() + _finalize(session, job_id=scope.job_ref, chunks=coarse_chunks) + session.commit() + + # Coarse retrieval works across the full document right after Phase A. + async def _discover_and_assemble(query_text: str) -> list[dict[str, object]]: + async with get_db_context() as async_db: + discovered = await map_unit_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query=query_text, + top_k=5, + exclude_document_ids=[], + exclude_sections=[], + ) + fused = list(discovered.payload.get("fused_rows") or []) + return await assemble_retrieval_results( + db=async_db, + rows=fused, + exclude_document_ids=[], + exclude_sections=[], + ) + + coarse_hits = asyncio.run(_discover_and_assemble("alpha coarse unique")) + assert any( + row["chunk_id"] == "coarse-p1-2" + and row["content"] == "alpha coarse unique token" + and row["content_source"] == "raw" + for row in coarse_hits + ) + + version_before = asyncio.run( + get_cache_version( + user_id=scope.user_ref, + namespace=scope.namespace_ref, + ) + ) + + with Session(engine) as session: + doc = session.execute( + select(Document).where( + Document.current_job_result_id == scope.revision_ref + ) + ).scalar_one() + sections_before = { + sec.section_path: sec.section_id + for sec in session.scalars( + select(DocumentSection).where( + DocumentSection.job_result_id == scope.revision_ref + ) + ).all() + } + parent_sec_id_before = sections_before["SharedParent"] + sibling_sec_id_before = sections_before["SharedParent / SiblingLeafB"] + sibling_chunk_before = session.execute( + select(DocumentChunk).where( + DocumentChunk.job_result_id == scope.revision_ref, + DocumentChunk.chunk_id == "coarse-p3-4", + ) + ).scalar_one() + sibling_chunk_pk_before = sibling_chunk_before.id + + replace_scope_content( + session, + scope=DocumentPublicationScope( + user_id=scope.user_ref, + namespace=scope.namespace_ref, + document_id=doc.document_id, + job_result_id=scope.revision_ref, + source_file_name=scope.source_file_name, + ), + scope_id="p1-2", + chunks=[ + { + "chunk_id": "refined-a1", + "type": "page", + "content": "delta refined first leaf raw body", + "path": f"{scope.source_file_name}/SharedParent/RefinedA1", + "metadata": { + "scope_id": "p1-2", + "summary": "delta refined summary a1", + "page_nums": [1], + }, + "order": 0, + }, + { + "chunk_id": "refined-a2", + "type": "page", + "content": "epsilon refined second leaf raw body", + "path": f"{scope.source_file_name}/SharedParent/RefinedA2", + "metadata": { + "scope_id": "p1-2", + "summary": "epsilon refined summary a2", + "page_nums": [2], + }, + "order": 1, + }, + ], + section_summaries={ + "SharedParent": "This section covers: RefinedA1, RefinedA2", + "SharedParent / RefinedA1": "delta refined summary a1", + "SharedParent / RefinedA2": "epsilon refined summary a2", + }, + ) + session.commit() + + version_after = asyncio.run( + get_cache_version( + user_id=scope.user_ref, + namespace=scope.namespace_ref, + ) + ) + assert version_after == version_before + 1 + + with Session(engine) as session: + sections_after = { + sec.section_path: sec + for sec in session.scalars( + select(DocumentSection).where( + DocumentSection.job_result_id == scope.revision_ref + ) + ).all() + } + assert "SharedParent / OldLeafA" not in sections_after + assert sections_after["SharedParent"].section_id == parent_sec_id_before + assert ( + sections_after["SharedParent / SiblingLeafB"].section_id + == sibling_sec_id_before + ) + assert ( + sections_after["SharedParent / RefinedA1"].parent_section_id + == parent_sec_id_before + ) + assert ( + sections_after["SharedParent / RefinedA2"].parent_section_id + == parent_sec_id_before + ) + + chunks_after = list( + session.scalars( + select(DocumentChunk) + .where(DocumentChunk.job_result_id == scope.revision_ref) + .order_by(DocumentChunk.sort_order) + ).all() + ) + assert [c.chunk_id for c in chunks_after] == [ + "refined-a1", + "refined-a2", + "coarse-p3-4", + "coarse-p5-6", + ] + assert chunks_after[2].id == sibling_chunk_pk_before + assert chunks_after[0].chunk_metadata["enrichment_state"] == "refined" + assert chunks_after[1].chunk_metadata["enrichment_state"] == "refined" + assert chunks_after[2].chunk_metadata["enrichment_state"] == "coarse" + assert chunks_after[3].chunk_metadata["enrichment_state"] == "coarse" + + map_index = session.execute( + select(DocumentMapUnitIndex).where( + DocumentMapUnitIndex.job_result_id == scope.revision_ref + ) + ).scalar_one() + assert map_index.unit_count == 4 + + # Refined scope now serves summary, while unrefined sibling scope still serves raw. + refined_hits = asyncio.run(_discover_and_assemble("delta refined first")) + assert any( + row["chunk_id"] == "refined-a1" + and row["content"] == "delta refined summary a1" + and row["content_source"] == "summary" + for row in refined_hits + ) + sibling_hits = asyncio.run(_discover_and_assemble("bravo sibling untouched")) + assert any( + row["chunk_id"] == "coarse-p3-4" + and row["content"] == "bravo sibling untouched token" + and row["content_source"] == "raw" + for row in sibling_hits + ) + finally: + engine.dispose() + + +def test_read_promotes_coarse_and_queued_scopes_with_dedup_including_descendants( + monkeypatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + from shared.core import celery_app as celery_app_mod + from shared.core.database import get_db_context + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("read-promote") + sent_tasks: list[tuple[str, list[object], str | None, int | None]] = [] + + class _FakeCelery: + def send_task( + self, + name: str, + args: list[object] | None = None, + queue: str | None = None, + priority: int | None = None, + **kwargs: object, + ) -> None: + del kwargs + sent_tasks.append((name, list(args or []), queue, priority)) + + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: _FakeCelery()) + + chunks = [ + _page_chunk( + chunk_id="chunk-p1-2", + path=f"{scope.source_file_name}/Chapter1/SecA", + scope_id="p1-2", + page_nums=[1, 2], + ), + _page_chunk( + chunk_id="chunk-p3-4", + path=f"{scope.source_file_name}/Chapter1/SecB", + scope_id="p3-4", + page_nums=[3, 4], + ), + { + "chunk_id": "chunk-p5-6", + "type": "page", + "content": "already refined", + "path": f"{scope.source_file_name}/Chapter1/SecC", + "metadata": { + "scope_id": "p5-6", + "enrichment_state": "refined", + "summary": "refined sec c", + "page_nums": [5, 6], + }, + }, + ] + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + _set_job_progressive(session, job_id=scope.job_ref) + session.commit() + _finalize(session, job_id=scope.job_ref, chunks=chunks) + _write_scope_rows(session, scope=scope) + # Mark p3-4 as low-queued and p5-6 as already refined before read. + for row in session.scalars( + select(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == scope.revision_ref + ) + ).all(): + if row.scope_id == "p3-4": + row.state = "queued" + row.priority = "low" + elif row.scope_id == "p5-6": + row.state = "refined" + doc_id = session.execute( + select(Document.document_id).where( + Document.current_job_result_id == scope.revision_ref + ) + ).scalar_one() + session.commit() + + async def _call_read() -> None: + async with get_db_context() as async_db: + ctx = ToolContext( + db=async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + db_factory=get_db_context, + ) + result = await read( + ctx, + { + "refs": [ + { + "document_id": doc_id, + "section_path": "Chapter1", + } + ], + "mode": "descendants", + }, + ) + assert result.error is None + + asyncio.run(_call_read()) + + assert sorted(sent_tasks, key=lambda item: str(item[1][1])) == [ + ( + ENRICH_SCOPE_TASK_NAME, + [scope.revision_ref, "p1-2"], + HIGH_ENRICHMENT_QUEUE, + HIGH_ENRICHMENT_TASK_PRIORITY, + ), + ( + ENRICH_SCOPE_TASK_NAME, + [scope.revision_ref, "p3-4"], + HIGH_ENRICHMENT_QUEUE, + HIGH_ENRICHMENT_TASK_PRIORITY, + ), + ] + + # Reading the same coarse/queued scopes a second time must not re-enqueue. + asyncio.run(_call_read()) + assert len(sent_tasks) == 2 + + with Session(engine) as session: + rows = ( + session.execute( + select(DocumentScopeEnrichment) + .where( + DocumentScopeEnrichment.job_result_id == scope.revision_ref + ) + .order_by(DocumentScopeEnrichment.scope_id) + ) + .scalars() + .all() + ) + assert [(r.scope_id, r.state, r.priority) for r in rows] == [ + ("p1-2", "queued", "high"), + ("p3-4", "queued", "high"), + ("p5-6", "refined", None), + ] + assert rows[0].last_hit_at is not None + assert rows[1].last_hit_at is not None + finally: + engine.dispose() + + +def test_full_mode_does_not_write_scope_rows_or_enqueue_tasks( + monkeypatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("full-mode") + chunks = [ + { + "chunk_id": "full-p1", + "type": "page", + "content": "full refined body", + "path": f"{scope.source_file_name}/A", + "metadata": { + "scope_id": "p1-2", + "enrichment_state": "coarse", + "summary": "full summary", + "page_nums": [1, 2], + }, + } + ] + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + _set_job_progressive( + session, + job_id=scope.job_ref, + ingestion_mode="full", + enrichment_status="", + ) + session.commit() + finalization = _finalize( + session, job_id=scope.job_ref, chunks=chunks + ) + session.commit() + + assert finalization.post_commit_effects.structure_job_result_ids == () + with Session(engine) as session: + rows = ( + session.execute( + select(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == scope.revision_ref + ) + ) + .scalars() + .all() + ) + assert rows == [] + finally: + engine.dispose() + + +def test_m0_text_publish_enqueues_structure_and_supports_any_stage_explore_and_ancestor_summary_isolation( + monkeypatch, + postgresql_proc: PostgreSQLProcess, +) -> None: + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + from shared.core import celery_app as celery_app_mod + from shared.core.database import get_db_context + from shared.services.jobs.lifecycle.post_commit_effects import ( + SyncJobPostCommitEffectRunner as PostCommitRunner, + ) + from shared.services.jobs.lifecycle.scope_enrichments import ( + write_coarse_scope_enrichments, + ) + from shared.services.retrieval.publication_scope_content import ( + replace_document_scope_content as replace_scope_content, + ) + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("m0-flow") + sent_tasks: list[tuple[str, list[object], str | None, int | None]] = [] + + class _FakeCelery: + def send_task( + self, + name: str, + args: list[object] | None = None, + queue: str | None = None, + priority: int | None = None, + **kwargs: object, + ) -> None: + del kwargs + sent_tasks.append((name, list(args or []), queue, priority)) + + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: _FakeCelery()) + + m0_chunks = [ + { + "chunk_id": "m0-p1", + "type": "page", + "content": "zeta baseline m0 page one unique text", + "path": f"{scope.source_file_name}/p1", + "metadata": { + "scope_id": "p1-1", + "enrichment_state": "text", + "page_nums": [1], + }, + }, + { + "chunk_id": "m0-p2", + "type": "page", + "content": "theta baseline m0 page two unique text", + "path": f"{scope.source_file_name}/p2", + "metadata": { + "scope_id": "p2-2", + "enrichment_state": "text", + "page_nums": [2], + }, + }, + { + "chunk_id": "m0-p3", + "type": "page", + "content": "kappa baseline m0 page three unique text", + "path": f"{scope.source_file_name}/p3", + "metadata": { + "scope_id": "p3-3", + "enrichment_state": "text", + "page_nums": [3], + }, + }, + ] + + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + _set_job_progressive( + session, + job_id=scope.job_ref, + ingestion_mode="progressive", + enrichment_status="text", + ) + session.commit() + finalization = _finalize( + session, job_id=scope.job_ref, chunks=m0_chunks + ) + doc_id = session.execute( + select(Document.document_id).where( + Document.current_job_result_id == scope.revision_ref + ) + ).scalar_one() + session.commit() + + assert finalization.post_commit_effects.structure_job_result_ids == ( + scope.revision_ref, + ) + PostCommitRunner().run(finalization.post_commit_effects) + assert sent_tasks == [ + ( + BUILD_STRUCTURE_TASK_NAME, + [scope.revision_ref], + NORMAL_ENRICHMENT_QUEUE, + None, + ) + ] + + # Stage M0 (state=text): recall/grep/read work; outline/node_filter/assets/query_table return informative notes. + async def _verify_m0_tools() -> None: + async with get_db_context() as async_db: + ctx = ToolContext( + db=async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + db_factory=get_db_context, + ) + recall_res = await recall(ctx, {"query": "zeta baseline m0"}) + assert recall_res.error is None + assert "candidates=1" not in recall_res.text or "zeta" in recall_res.text + assert len(recall_res.payload["rows"]) >= 1 + + grep_res = await grep(ctx, {"pattern": "zeta baseline"}) + assert grep_res.error is None + assert "zeta baseline" in grep_res.text + + read_res = await read( + ctx, + {"refs": [{"document_id": doc_id, "chunk_id": "m0-p1"}]}, + ) + assert read_res.error is None + assert "zeta baseline m0 page one unique text" in read_res.text + + outline_res = await outline( + ctx, {"scope": [{"document_id": doc_id}]} + ) + assert outline_res.error is None + assert "still being structured (state=text)" in outline_res.text + + nf_res = await node_filter( + ctx, + { + "scope": [{"document_id": doc_id}], + "predicates": [{"field": "path", "terms": ["Root"]}], + }, + ) + assert nf_res.error is None + assert "still being structured (state=text)" in nf_res.text + + assets_res = await assets( + ctx, {"scope": [{"document_id": doc_id}]} + ) + assert assets_res.error is None + assert "still being structured (state=text)" in assets_res.text + + qt_res = await query_table( + ctx, + { + "document_id": doc_id, + "chunk_id": "nonexistent_tbl", + "sql": "SELECT * FROM t", + }, + ) + assert qt_res.error is not None + assert "still being structured (state=text)" in qt_res.error + + ld_res = await list_documents(ctx, {}) + assert ld_res.error is None + assert "status=text" in ld_res.text + assert ld_res.payload["documents"][0]["status"] == "text" + + asyncio.run(_verify_m0_tools()) + # Reading an M0 page must NOT enqueue scope tasks (no acceleration channel). + assert len(sent_tasks) == 1 + + # Stage 1 (coarse): whole-doc restructure publish + write scope rows. + from shared.models.database.job_result import JobResult + + with Session(engine) as session: + _set_job_progressive( + session, + job_id=scope.job_ref, + ingestion_mode="progressive", + enrichment_status="coarse", + ) + jr = session.execute( + select(JobResult).where(JobResult.id == scope.revision_ref) + ).scalar_one() + jr.document_metadata = { + **dict(jr.document_metadata or {}), + "enrichment_status": "coarse", + } + pub_scope = DocumentPublicationScope( + user_id=scope.user_ref, + namespace=scope.namespace_ref, + document_id=doc_id, + job_result_id=scope.revision_ref, + source_file_name=scope.source_file_name, + ) + replace_scope_content( + session, + scope=pub_scope, + scope_id="*", + chunks=[ + _page_chunk( + chunk_id="coarse-s1", + path=f"{scope.source_file_name}/Chapter1/SecA", + scope_id="p1-1", + page_nums=[1], + ), + _page_chunk( + chunk_id="coarse-s2", + path=f"{scope.source_file_name}/Chapter1/SecB", + scope_id="p2-2", + page_nums=[2], + ), + ], + section_summaries={ + "Chapter1": "Original full Chapter1 summary covering SecA and SecB", + "Chapter1 / SecA": "SecA coarse summary", + "Chapter1 / SecB": "SecB coarse summary", + }, + ) + write_coarse_scope_enrichments( + session, + document_id=doc_id, + job_result_id=scope.revision_ref, + ) + session.commit() + + sent_tasks.clear() + + async def _verify_coarse_tools() -> None: + async with get_db_context() as async_db: + ctx = ToolContext( + db=async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + db_factory=get_db_context, + ) + outline_res = await outline( + ctx, {"scope": [{"document_id": doc_id}]} + ) + assert outline_res.error is None + assert "Chapter1 / SecA" in outline_res.text + + assets_res = await assets( + ctx, + { + "scope": [ + { + "document_id": doc_id, + "section_path": "Chapter1 / SecA", + } + ] + }, + ) + assert assets_res.error is None + assert "not extracted yet (state=coarse)" in assets_res.text + + ld_res = await list_documents(ctx, {}) + assert ld_res.payload["documents"][0]["status"] == "coarse" + + asyncio.run(_verify_coarse_tools()) + # Calling assets on a coarse scope promotes that scope to high priority! + assert sent_tasks == [ + ( + ENRICH_SCOPE_TASK_NAME, + [scope.revision_ref, "p1-1"], + HIGH_ENRICHMENT_QUEUE, + HIGH_ENRICHMENT_TASK_PRIORITY, + ) + ] + + # Stage 2 (fine) on p1-1: verify shared ancestor 'Chapter1' summary is NOT overwritten. + with Session(engine) as session: + replace_scope_content( + session, + scope=DocumentPublicationScope( + user_id=scope.user_ref, + namespace=scope.namespace_ref, + document_id=doc_id, + job_result_id=scope.revision_ref, + source_file_name=scope.source_file_name, + ), + scope_id="p1-1", + chunks=[ + { + "chunk_id": "fine-s1", + "type": "page", + "content": "refined sec a", + "path": f"{scope.source_file_name}/Chapter1/SecA", + "metadata": { + "scope_id": "p1-1", + "summary": "SecA fine summary", + "page_nums": [1], + }, + } + ], + section_summaries={ + "Chapter1": "Partial overwritten summary mentioning only SecA", + "Chapter1 / SecA": "SecA fine summary", + }, + ) + session.commit() + ch1 = session.execute( + select(DocumentSection).where( + DocumentSection.job_result_id == scope.revision_ref, + DocumentSection.section_path == "Chapter1", + ) + ).scalar_one() + assert ( + ch1.summary + == "Original full Chapter1 summary covering SecA and SecB" + ) + finally: + engine.dispose() + + +def test_scanned_visual_m0_retrieval_evidence_and_cross_stage_embedding_reuse( + monkeypatch, + postgresql_proc: PostgreSQLProcess, + tmp_path, +) -> None: + import base64 + from pathlib import Path + + configure_contract_environment(monkeypatch, postgresql_proc) + asyncio.run(prepare_contract_storage()) + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + scope = _scope("vism0") + + from shared.core.config import settings + from shared.core.database import get_db_context + from shared.models.database.document import DocumentPageEmbedding + from shared.services.ai.embedding import DEFAULT_EMBEDDING_DIM, multimodal_client + from shared.services.retrieval.agent_tools.tools.grep import grep as fresh_grep + from shared.services.retrieval.agent_tools.tools.read import read as fresh_read + from shared.services.retrieval.agent_tools.tools.recall import recall as fresh_recall + from shared.services.retrieval.execution.route_types import RetrievalRouteContext + from shared.services.retrieval.execution.routes import _run_classic_topk_route + from shared.services.retrieval.publication_scope_content import ( + replace_document_scope_content as replace_scope_content, + ) + from shared.services.retrieval.search.visual_discovery import ( + hybrid_discovery, + visual_page_discovery, + ) + from shared.services.storage.result_storage import JobResultStorage + + vec_p1 = [1.0] + [0.0] * (DEFAULT_EMBEDDING_DIM - 1) + vec_p2 = [0.0, 1.0] + [0.0] * (DEFAULT_EMBEDDING_DIM - 2) + monkeypatch.setattr(settings, "VISUAL_RETRIEVAL_ENABLED", True) + monkeypatch.setattr( + multimodal_client, + "_call_gateway_embedding", + lambda *, text=None, **kw: ( + vec_p2 + if "second page chart" in str(text or "") + else vec_p1 + ), + ) + + fake_png_bytes = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" + ) + fake_png_file = Path(tmp_path) / "page.png" + fake_png_file.write_bytes(fake_png_bytes) + + monkeypatch.setattr( + JobResultStorage, + "download_raw_to_temp", + lambda self, **kwargs: str(fake_png_file), + ) + monkeypatch.setattr( + JobResultStorage, + "generate_artifact_url", + lambda self, **kwargs: f"https://assets.knowhere.test/{kwargs.get('artifact_ref', 'page.png')}", + ) + + try: + with Session(engine) as session: + seed_publication_fixture(session, scope=scope) + _set_job_progressive( + session, + job_id=scope.job_ref, + ingestion_mode="progressive", + enrichment_status="text", + ) + session.commit() + + # Publish pure-visual M0 (empty content, _page_embedding + page_assets). + m0_chunks: list[dict[str, object]] = [ + { + "chunk_id": "scan-m0-p1", + "type": "page", + "content": "", + "path": f"{scope.source_file_name}/p1", + "metadata": { + "scope_id": "p1-1", + "enrichment_state": "text", + "page_nums": [1], + "_page_embedding": vec_p1, + "page_assets": [ + { + "page_num": 1, + "artifact_ref": "page_citation_assets/page-1.png", + "content_type": "image/png", + "source": "knowhere-rendered-page-citation-source", + } + ], + }, + }, + { + "chunk_id": "scan-m0-p2", + "type": "page", + "content": "", + "path": f"{scope.source_file_name}/p2", + "metadata": { + "scope_id": "p2-2", + "enrichment_state": "text", + "page_nums": [2], + "_page_embedding": vec_p2, + "page_assets": [ + { + "page_num": 2, + "artifact_ref": "page_citation_assets/page-2.png", + "content_type": "image/png", + "source": "knowhere-rendered-page-citation-source", + } + ], + }, + }, + ] + with Session(engine) as session: + _finalize(session, job_id=scope.job_ref, chunks=m0_chunks) + doc_id = str( + session.execute( + select(Document.document_id).where( + Document.current_job_result_id == scope.revision_ref + ) + ).scalar_one() + ) + session.commit() + + emb_rows = list( + session.scalars( + select(DocumentPageEmbedding) + .where(DocumentPageEmbedding.document_id == doc_id) + .order_by(DocumentPageEmbedding.page_num) + ).all() + ) + assert len(emb_rows) == 2 + assert [r.chunk_id for r in emb_rows] == ["scan-m0-p1", "scan-m0-p2"] + assert [r.page_num for r in emb_rows] == [1, 2] + + persisted_chunks = list( + session.scalars( + select(DocumentChunk) + .where(DocumentChunk.document_id == doc_id) + .order_by(DocumentChunk.sort_order) + ).all() + ) + assert all( + "_page_embedding" not in (c.chunk_metadata or {}) + for c in persisted_chunks + ) + + async def _verify_visual_m0_retrieval_and_tools() -> None: + async with get_db_context() as async_db: + # 1. Direct visual_page_discovery hits page 2 for query matching vec_p2. + vis_res = await visual_page_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="what does the second page chart show", + top_k=2, + ) + vis_rows = vis_res.payload["fused_rows"] + assert len(vis_rows) == 2 + assert vis_rows[0]["chunk_id"] == "scan-m0-p2" + + # 2. Classic route answers via visual channel even though lexical is empty. + classic_outcome = await _run_classic_topk_route( + RetrievalRouteContext( + db=async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="what does the second page chart show", + top_k=1, + effective_recall_k=2, + use_agentic=False, + exclude_document_ids=[], + exclude_sections=[], + chunk_types=None, + allowed_chunk_types=None, + rerank=False, + threshold=0.0, + internal_recall_k=2, + signal_paths=[], + filter_mode="delete", + ) + ) + classic_results = classic_outcome.response["results"] + assert len(classic_results) == 1 + assert classic_results[0]["chunk_id"] == "scan-m0-p2" + assert classic_results[0]["content_source"] == "page_image" + assert classic_results[0]["page_nums"] == [2] + assert any( + part.get("type") == "image" + for part in classic_results[0]["composed"] + ) + + # 3. Agent corpus.recall + corpus.read work on pure-visual M0. + ctx = ToolContext( + db=async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + db_factory=get_db_context, + ) + recall_res = await fresh_recall( + ctx, {"query": "what does the second page chart show"} + ) + assert recall_res.error is None + assert recall_res.payload["rows"][0]["section_path"] == "p2" + assert "scanned page image" in recall_res.text + + grep_res = await fresh_grep( + ctx, {"pattern": "chart"} + ) + assert grep_res.error is None + assert "scanned document still awaiting background OCR" in grep_res.text + assert "corpus.recall" in grep_res.text + + read_res = await fresh_read( + ctx, + { + "refs": [ + {"document_id": doc_id, "chunk_id": "scan-m0-p2"} + ] + }, + ) + assert read_res.error is None + assert "[Page image: https://assets.knowhere.test/page_citation_assets/page-2.png]" in read_res.text + assert read_res.media == [ + { + "type": "image_url", + "url": "https://assets.knowhere.test/page_citation_assets/page-2.png", + } + ] + + asyncio.run(_verify_visual_m0_retrieval_and_tools()) + + # 4. Background coarse + fine republish reuse M0 embeddings without recomputing, + # rebind chunk_id to new section chunks, and inherit page_assets. + with Session(engine) as session: + pub_scope = DocumentPublicationScope( + user_id=scope.user_ref, + namespace=scope.namespace_ref, + document_id=doc_id, + job_result_id=scope.revision_ref, + source_file_name=scope.source_file_name, + ) + replace_scope_content( + session, + scope=pub_scope, + scope_id="*", + chunks=[ + { + "chunk_id": "coarse-sec-1", + "type": "page", + "content": "ocr text page one", + "path": f"{scope.source_file_name}/Chapter1/SecA", + "metadata": { + "scope_id": "p1-1", + "enrichment_state": "coarse", + "page_nums": [1], + }, + }, + { + "chunk_id": "coarse-sec-2", + "type": "page", + "content": "ocr text page two", + "path": f"{scope.source_file_name}/Chapter1/SecB", + "metadata": { + "scope_id": "p2-2", + "enrichment_state": "coarse", + "page_nums": [2], + }, + }, + ], + section_summaries={"Chapter1 / SecB": "SecB summary"}, + ) + session.commit() + + emb_rows_after_coarse = list( + session.scalars( + select(DocumentPageEmbedding) + .where(DocumentPageEmbedding.document_id == doc_id) + .order_by(DocumentPageEmbedding.page_num) + ).all() + ) + assert [r.chunk_id for r in emb_rows_after_coarse] == [ + "coarse-sec-1", + "coarse-sec-2", + ] + sec2_chunk = session.execute( + select(DocumentChunk).where( + DocumentChunk.document_id == doc_id, + DocumentChunk.chunk_id == "coarse-sec-2", + ) + ).scalar_one() + assert (sec2_chunk.chunk_metadata or {}).get("page_assets") + + async def _verify_hybrid_after_coarse() -> float: + async with get_db_context() as async_db: + hyb_res = await hybrid_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="second page chart", + top_k=2, + ) + assert hyb_res.payload["fused_rows"][0]["chunk_id"] == "coarse-sec-2" + + vis_res = await visual_page_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="second page chart", + top_k=2, + ) + vis_rows = vis_res.payload["fused_rows"] + assert vis_rows[0]["chunk_id"] == "coarse-sec-2" + + excluded = await visual_page_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="second page chart", + top_k=2, + exclude_sections=[ + { + "document_id": doc_id, + "section_path": "Chapter1 / SecB", + } + ], + ) + assert "coarse-sec-2" not in [ + row["chunk_id"] for row in excluded.payload["fused_rows"] + ] + return float(vis_rows[0]["score"]) + + page2_score = asyncio.run(_verify_hybrid_after_coarse()) + + with Session(engine) as session: + replace_scope_content( + session, + scope=pub_scope, + scope_id="*", + chunks=[ + { + "chunk_id": "merged-sec", + "type": "page", + "content": "ocr text page one and two", + "path": f"{scope.source_file_name}/Chapter1", + "metadata": { + "scope_id": "p1-1", + "enrichment_state": "coarse", + "page_nums": [1, 2], + }, + } + ], + ) + session.commit() + merged_emb_rows = list( + session.scalars( + select(DocumentPageEmbedding) + .where(DocumentPageEmbedding.document_id == doc_id) + .order_by(DocumentPageEmbedding.page_num) + ).all() + ) + assert [r.chunk_id for r in merged_emb_rows] == [ + "merged-sec", + "merged-sec", + ] + assert [r.page_num for r in merged_emb_rows] == [1, 2] + + async def _verify_chunk_max_pool_and_flag_off() -> None: + async with get_db_context() as async_db: + merged_res = await visual_page_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="second page chart", + top_k=2, + ) + merged_rows = merged_res.payload["fused_rows"] + assert len(merged_rows) == 1 + assert merged_rows[0]["chunk_id"] == "merged-sec" + assert abs(float(merged_rows[0]["score"]) - page2_score) < 1e-9 + + monkeypatch.setattr(settings, "VISUAL_RETRIEVAL_ENABLED", False) + off_res = await visual_page_discovery( + async_db, + user_id=scope.user_ref, + namespace=scope.namespace_ref, + query="second page chart", + top_k=2, + ) + assert off_res.payload["fused_rows"] == [] + + asyncio.run(_verify_chunk_max_pool_and_flag_off()) + finally: + engine.dispose() + + diff --git a/apps/worker/.env.example b/apps/worker/.env.example index fe58ed661..0200465ab 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -78,6 +78,7 @@ GPT_API_KEY= GLM_API_KEY= ALI_API_KEYS= ARK_API_KEY= +VERCEL_KEY= # Optional overrides: # DS_URL=https://api.deepseek.com/v1 @@ -91,6 +92,9 @@ ARK_API_KEY= # ASSET_MODEL is optional; empty follows IMAGE_MODEL. # IMAGE_MODEL=deepseek-flash # ASSET_MODEL= +# Page-image embeddings use VERCEL_KEY. +# EMBEDDING_MODEL=google/gemini-embedding-2 +# EMBEDDING_DIM=3072 # Optional retrieval overrides have code defaults. Retrieval is evidence-only: # evidence is the primary output; evidence_text is deprecated and always empty; diff --git a/apps/worker/app/core/tasks/base_task.py b/apps/worker/app/core/tasks/base_task.py index c3474884f..4142d9ed1 100644 --- a/apps/worker/app/core/tasks/base_task.py +++ b/apps/worker/app/core/tasks/base_task.py @@ -16,6 +16,8 @@ class DocumentIngestionBaseTask(Task): """Base task class for worker-side Document Ingestion error handling.""" + finalize_job_on_failure: bool = True + def on_success(self, retval, task_id, args, kwargs): """Task success callback.""" logger.bind( @@ -66,6 +68,8 @@ def on_failure(self, exc, task_id, args, kwargs, einfo): def _extract_job_id(self, args, kwargs) -> Optional[str]: """Extract job_id from args or kwargs.""" + if not self.finalize_job_on_failure: + return None if args and len(args) > 0: if isinstance(args[0], dict) and "job_id" in args[0]: return args[0]["job_id"] diff --git a/apps/worker/app/core/tasks/enrichment_tasks.py b/apps/worker/app/core/tasks/enrichment_tasks.py new file mode 100644 index 000000000..4600d5ee0 --- /dev/null +++ b/apps/worker/app/core/tasks/enrichment_tasks.py @@ -0,0 +1,90 @@ +"""Celery tasks for Phase B page-memory scope enrichment.""" + +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from app.core.tasks.base_task import DocumentIngestionBaseTask +from app.services.page_memory.scope_enrichment_service import ( + execute_scope_enrichment, + execute_structure_enrichment, +) +from shared.core.celery_app import get_celery_app +from shared.core.logging import LogEvent, log_context +from shared.services.jobs.lifecycle.scope_enrichments import ( + sweep_progressive_enrichments_sync, +) +from shared.services.redis.periodic_task_lock import periodic_task_lock + +celery_app = get_celery_app() +_SWEEP_PERIOD_SECONDS = 300 + + +class ScopeEnrichmentBaseTask(DocumentIngestionBaseTask): + """Scope enrichment task base that keeps the parent Job in DONE.""" + + finalize_job_on_failure: bool = False + + +@celery_app.task( + bind=True, + base=ScopeEnrichmentBaseTask, + name="app.core.tasks.enrichment_tasks.enrich_page_memory_scope", + ignore_result=True, +) +def enrich_page_memory_scope( + self: Any, + job_result_id: str, + scope_id: str, +) -> dict[str, Any]: + """Run Phase B C1-C7 refinement for one published page-memory scope.""" + with log_context(task_id=self.request.id): + result = execute_scope_enrichment( + job_result_id=str(job_result_id), + scope_id=str(scope_id), + ) + logger.bind(event=LogEvent.WORKER_TASK_COMPLETE.value).info( + "Task completed: enrich_page_memory_scope " + f"job_result_id={job_result_id} scope_id={scope_id} " + f"status={result.get('status')}" + ) + return result + + +@celery_app.task( + bind=True, + base=ScopeEnrichmentBaseTask, + name="app.core.tasks.enrichment_tasks.build_page_memory_structure", + ignore_result=True, +) +def build_page_memory_structure( + self: Any, + job_result_id: str, +) -> dict[str, Any]: + """Run background M0 -> coarse structure build for a progressive page-memory revision.""" + with log_context(task_id=self.request.id): + result = execute_structure_enrichment( + job_result_id=str(job_result_id), + ) + logger.bind(event=LogEvent.WORKER_TASK_COMPLETE.value).info( + "Task completed: build_page_memory_structure " + f"job_result_id={job_result_id} status={result.get('status')}" + ) + return result + + +@celery_app.task( + name="app.core.tasks.enrichment_tasks.sweep_progressive_enrichments", + ignore_result=True, +) +def sweep_progressive_enrichments() -> dict[str, Any]: + """Periodic sweep: re-enqueue stuck M0 text revisions and retry failed scopes.""" + with periodic_task_lock( + "app.core.tasks.enrichment_tasks.sweep_progressive_enrichments", + period_seconds=_SWEEP_PERIOD_SECONDS, + ) as acquired: + if not acquired: + return {"status": "skipped", "reason": "duplicate Beat firing"} + return sweep_progressive_enrichments_sync() diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 2d1f8b411..1faac8cad 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -19,6 +19,7 @@ def _register_task_modules() -> None: """Import task modules for Celery side-effect registration.""" import app.core.tasks.document_ingestion_tasks # noqa: F401 + import app.core.tasks.enrichment_tasks # noqa: F401 import app.core.tasks.stale_job_sweeper # noqa: F401 import app.core.tasks.webhook_tasks # noqa: F401 diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 7cb899279..ee3f82f44 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -300,6 +300,10 @@ def _run_text_scan(self) -> None: self.blackboard.page_text_search_view = None return if profile.is_scanned: + if self.ctx.settings.get("skip_scanned_ocr"): + self.blackboard.page_full_text_cache = {} + self.blackboard.page_text_search_view = None + return result = REGISTRY.dispatch("ocr.pages", self.ctx, {"pages": pages}) self.trace.record_step( round_index=self.round_index, diff --git a/apps/worker/app/services/document_ingestion/parse_execution.py b/apps/worker/app/services/document_ingestion/parse_execution.py index 19edf519c..49e81762b 100644 --- a/apps/worker/app/services/document_ingestion/parse_execution.py +++ b/apps/worker/app/services/document_ingestion/parse_execution.py @@ -133,6 +133,9 @@ def _execute_page_memory_parse( page_memory_config = JobMetadataHelper.get_page_memory_config( job_context.job_metadata, ) + ingestion_mode = JobMetadataHelper.get_ingestion_mode( + job_context.job_metadata, + ) page_output_dir, parsed_df = run( PageMemoryInput( file_path=prepared_source.local_file_path, @@ -141,6 +144,7 @@ def _execute_page_memory_parse( output_dir=output_dir, job_id=job_id, page_memory_config=page_memory_config, + ingestion_mode=ingestion_mode, ) ) return ParseOutput(output_dir=page_output_dir, parsed_df=parsed_df) diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py index 0facf5a52..b0d473f1b 100644 --- a/apps/worker/app/services/document_ingestion/success_finalization.py +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -28,6 +28,10 @@ from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.ai.token_tracking import get_current_token_tracker from app.services.document_parser.support.stage_profiler import get_current_stage_tracker, stage_timer +from app.services.page_memory.phase_a import ( + PHASE_B_BLACKBOARD_REL, + existing_phase_b_input_refs, +) from shared.services.storage.result_storage import ResultStorage, get_result_storage from shared.services.storage.zip_result_service import ZipResultService @@ -261,6 +265,19 @@ def _upload_result_package( add_dir = str(result_package.artifact.add_dir) if result_package.artifact.add_dir else "" if add_dir and os.path.isfile(os.path.join(add_dir, "source.pdf")): artifact_refs.add("source.pdf") + parse_track = ( + JobMetadataHelper.get_parse_track(job_context.job_metadata) + if job_context is not None + else "" + ) + ingestion_mode = ( + JobMetadataHelper.get_ingestion_mode(job_context.job_metadata) + if job_context is not None + else "full" + ) + phase_b_refs = existing_phase_b_input_refs(add_dir) if add_dir else [] + if parse_track == "page_memory" and ingestion_mode == "progressive": + artifact_refs.update(phase_b_refs) with logfire.span("Upload result assets for {job_id}", job_id=job_id), stage_timer("worker.result.assets_upload", job_id=job_id): result_bundle = result_storage_factory().upload( job_id=job_id, @@ -268,6 +285,23 @@ def _upload_result_package( zip_file_path=generated_package.zip_file_path, artifact_refs=artifact_refs, ) + metadata_updates: dict[str, object] = {} + if job_context is not None and parse_track == "page_memory": + metadata_updates["result_raw_prefix"] = result_bundle.raw_prefix + if ( + ingestion_mode == "progressive" + and add_dir + and os.path.isfile(os.path.join(add_dir, PHASE_B_BLACKBOARD_REL)) + ): + metadata_updates["page_memory_phase_b"] = phase_b_refs + metadata_updates["enrichment_status"] = "text" if job_context is not None and job_context.job_metadata.get("corpus_target") == "DEMO": - persist_job_metadata_updates(job_id=job_id, job_context=job_context, metadata_updates={"demo_asset_manifest": sorted(result_bundle.raw_files), "result_raw_prefix": result_bundle.raw_prefix}) + metadata_updates["demo_asset_manifest"] = sorted(result_bundle.raw_files) + metadata_updates.setdefault("result_raw_prefix", result_bundle.raw_prefix) + if job_context is not None and metadata_updates: + persist_job_metadata_updates( + job_id=job_id, + job_context=job_context, + metadata_updates=metadata_updates, + ) return result_bundle.zip_key diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index f591940ea..af57a8ab6 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -36,6 +36,8 @@ def profile_document( skip_shard_plan: bool = False, oversized_policy: Literal["chunk", "page_memory"] = "chunk", skip_toc_anchoring: bool = False, + stop_after: Literal["asset_probe"] | None = None, + skip_scanned_ocr: bool = False, ) -> ParserDocumentProfile: """ General document profiling entry point. @@ -75,6 +77,8 @@ def profile_document( skip_shard_plan=skip_shard_plan, oversized_policy=oversized_policy, skip_toc_anchoring=skip_toc_anchoring, + stop_after=stop_after, + skip_scanned_ocr=skip_scanned_ocr, ) finally: if not visual_debug_enabled(): @@ -97,6 +101,8 @@ def _profile_pdf( skip_shard_plan: bool = False, oversized_policy: Literal["chunk", "page_memory"] = "chunk", skip_toc_anchoring: bool = False, + stop_after: Literal["asset_probe"] | None = None, + skip_scanned_ocr: bool = False, ) -> ParserDocumentProfile: with _profile_db_context(enabled=bool(job_id)) as db: return _profile_pdf_with_db( @@ -108,6 +114,8 @@ def _profile_pdf( skip_shard_plan=skip_shard_plan, oversized_policy=oversized_policy, skip_toc_anchoring=skip_toc_anchoring, + stop_after=stop_after, + skip_scanned_ocr=skip_scanned_ocr, ) @@ -121,6 +129,8 @@ def _profile_pdf_with_db( skip_shard_plan: bool = False, oversized_policy: Literal["chunk", "page_memory"] = "chunk", skip_toc_anchoring: bool = False, + stop_after: Literal["asset_probe"] | None = None, + skip_scanned_ocr: bool = False, ) -> ParserDocumentProfile: profile_job_id = job_id or filename agent_output_dir = os.path.join(output_dir, "_doc_agent") if output_dir else None @@ -140,6 +150,8 @@ def _profile_pdf_with_db( "vlm_model": settings.IMAGE_MODEL, "toc_profile_enabled": page_toc_enabled, "skip_toc_anchoring": bool(skip_toc_anchoring), + "stop_after_asset_probe": stop_after == "asset_probe", + "skip_scanned_ocr": bool(skip_scanned_ocr), "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, }, ) @@ -168,6 +180,12 @@ def _profile_pdf_with_db( coordinator.blackboard.page_full_text_cache ), ) + setattr(profile, "blackboard", coordinator.blackboard) + if stop_after == "asset_probe": + if trace := getattr(coordinator, "trace", None): + trace.persist_doc_profile(profile) + setattr(profile, "trace_recorder", trace) + return profile if profile.page_count > settings.MAX_PDF_PAGE_LIMIT: if oversized_policy != "page_memory": raise_if_oversized_pdf_not_supported(page_count=profile.page_count) diff --git a/apps/worker/app/services/page_memory/_utils.py b/apps/worker/app/services/page_memory/_utils.py index 2635f2c00..10fe9992e 100644 --- a/apps/worker/app/services/page_memory/_utils.py +++ b/apps/worker/app/services/page_memory/_utils.py @@ -21,21 +21,6 @@ def is_fallback_root(skeleton: Any) -> bool: ) -def derive_excluded_toc_pages( - *, - skeletons: list[Any], - toc_pages: list[int] | None, -) -> list[int]: - """Exclude chapter TOCs while retaining every physical prefix page.""" - prefix_pages: set[int] = { - page - for skeleton in skeletons - if is_document_prefix(skeleton) - for page in range(skeleton.start_page, skeleton.end_page + 1) - } - return sorted({int(page) for page in (toc_pages or [])} - prefix_pages) - - def sort_skeletons(skeletons: list[Any]) -> list[Any]: """Order skeletons by start page only. diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index fc3c3ad5b..1288cae00 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -20,9 +20,7 @@ from app.services.page_memory.normalizer import normalize_to_pdf from app.services.page_memory._utils import ( CoarseScope, - build_hierarchy_scopes, collapse_page_ranges, - derive_excluded_toc_pages, page_scope_info, sort_skeletons, ) @@ -32,9 +30,16 @@ write_scope_artifacts as _write_scope_artifacts, write_top_level_artifacts as _write_top_level_artifacts, ) +from app.services.page_memory.phase_a import ( + build_coarse_text_dataframe, + build_phase_a, + build_visual_page_dataframe, + persist_m0_inputs, +) from loguru import logger +from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException from shared.models.schemas.page_memory_config import PageMemoryConfig, resolve_asset_model from app.services.page_memory.body_text_preview import build_body_text_preview @@ -51,6 +56,7 @@ class PageMemoryInput: page_memory_config: PageMemoryConfig = field( default_factory=PageMemoryConfig.default, ) + ingestion_mode: str = "full" @dataclass(frozen=True) @@ -67,7 +73,7 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: """Run the page-memory track. Supports two granularity verdicts: - - ``whole_doc`` (≤6 pages, no TOC) → single whole-document chunk + - ``whole_doc`` (≤5 pages, no TOC) → single whole-document chunk - ``page`` → leaf section nodes via the full C1-C7 pipeline """ full_output_dir = _resolve_output_dir(request) @@ -86,6 +92,10 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: base_url=request.base_url, ) _persist_source_pdf(pdf_path=pdf_path, output_dir=full_output_dir) + is_progressive = str(request.ingestion_mode or "full").strip().lower() == "progressive" + visual_m0_enabled = is_progressive and bool( + getattr(settings, "VISUAL_RETRIEVAL_ENABLED", True) + ) with stage_timer("page_memory.profile", filename=pdf_filename): profile = profile_document( pdf_path, @@ -94,6 +104,8 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: output_dir=full_output_dir, skip_shard_plan=True, oversized_policy="page_memory", + stop_after="asset_probe" if is_progressive else None, + skip_scanned_ocr=visual_m0_enabled, ) trace_recorder = getattr(profile, "trace_recorder", None) page_count = max(int(profile.page_count or 0), 0) @@ -108,10 +120,11 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: "description": f"PDF has {page_count} pages, limit is {page_memory_config.max_pages}", }], ) - verdict = _decide_granularity(profile) + verdict = "page" if is_progressive else _decide_granularity(profile) trace_summary = { "page_count": page_count, "verdict": verdict, + "ingestion_mode": "progressive" if is_progressive else "full", } if verdict == "whole_doc": @@ -125,6 +138,29 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: ), body_text_by_page=getattr(profile, "page_body_text_cache", None), ) + elif is_progressive: + if visual_m0_enabled and bool(getattr(profile, "is_scanned", False)): + parsed_df = _build_visual_m0_dataframe( + pdf_path=pdf_path, + filename=request.filename, + output_dir=full_output_dir, + page_count=page_count, + ) + else: + parsed_df = build_coarse_text_dataframe( + page_texts=coerce_page_text_cache( + getattr(profile, "page_full_text_cache", None) + ), + body_text_by_page=coerce_page_text_cache( + getattr(profile, "page_body_text_cache", None) + ), + page_count=page_count, + filename=request.filename, + ) + persist_m0_inputs( + output_dir=full_output_dir, + profile=profile, + ) else: parsed_df = _build_page_dataframe( pdf_path=pdf_path, @@ -132,8 +168,8 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: output_dir=full_output_dir, profile=profile, verdict=verdict, - trace_recorder=trace_recorder, page_memory_config=page_memory_config, + trace_recorder=trace_recorder, ) final_status = "success" trace_summary["rows_count"] = len(parsed_df) @@ -165,6 +201,62 @@ def _persist_source_pdf(*, pdf_path: str, output_dir: str) -> str: return source_pdf_path +def _build_visual_m0_dataframe( + *, + pdf_path: str, + filename: str, + output_dir: str, + page_count: int, +) -> pd.DataFrame: + from app.services.page_memory.page_renderer import render_document_pages + from shared.services.ai.embedding import embed_images + + pages = list(range(1, max(int(page_count or 0), 0) + 1)) + with stage_timer("page_memory.m0_render", filename=filename, page_count=len(pages)): + rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + scope_id="m0", + pages=pages, + page_texts={}, + ) + image_path_by_page = { + int(item.page_index): str(item.image_path) + for item in rendered + if item.image_path + } + ordered_pages = [page for page in pages if page in image_path_by_page] + with stage_timer( + "page_memory.m0_embed", + filename=filename, + page_count=len(ordered_pages), + ): + vectors = embed_images( + [image_path_by_page[page] for page in ordered_pages], + usage_task="page_memory.m0_embed", + ) + embeddings_by_page = {} + skipped_pages: list[int] = [] + for page, vec in zip(ordered_pages, vectors, strict=True): + if isinstance(vec, list) and vec: + embeddings_by_page[page] = vec + else: + skipped_pages.append(page) + if skipped_pages: + logger.warning( + "Skipped page embeddings after retries; document continues: pages={}", + skipped_pages, + ) + return build_visual_page_dataframe( + page_count=page_count, + filename=filename, + image_path_by_page=image_path_by_page, + embeddings_by_page=embeddings_by_page, + output_dir=output_dir, + ) + + def _decide_granularity(profile: Any) -> str: """Decide granularity for a page-memory document. @@ -177,21 +269,11 @@ def _decide_granularity(profile: Any) -> str: page_count = int(getattr(profile, "page_count", 0) or 0) toc = getattr(profile, "toc", None) has_toc = bool(getattr(toc, "has_toc", False)) - if page_count <= 6 and not has_toc: + if page_count <= 5 and not has_toc: return "whole_doc" return "page" -def _resolve_asset_max_pages( - page_count: int, - page_memory_config: PageMemoryConfig, -) -> int: - configured_limit = page_memory_config.asset_max_pages - if configured_limit is None: - return page_count - return max(0, min(page_count, configured_limit)) - - def _resolve_hierarchy_model(page_memory_config: PageMemoryConfig) -> str | None: return ( page_memory_config.hierarchy_model @@ -203,6 +285,70 @@ def _resolve_hierarchy_model(page_memory_config: PageMemoryConfig) -> str | None # ── page builder (C1→C2→C3→C4→C6→C7) ──────────────────────────────── +def assemble_c7_node_rows( + *, + skeletons: list[Any], + rendered: list[Any] | dict[int, Any], + tags: list[Any], + assets_by_page: dict[int, list[Any]], + page_texts: dict[int, str], + body_text_by_page: dict[int, str] | None, + page_labels: list[Any], + toc_pages: list[int], + page_count: int, + output_dir: str, + filename: str, + verdict: str, + vlm_model: str | None, + page_memory_config: PageMemoryConfig, +) -> tuple[list[dict[str, Any]], list[int]]: + """Shared C7 node-row assembly used by whole-doc ``_build_page_dataframe`` and per-scope refinement.""" + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + from app.services.page_memory.node_assembler import build_node_rows + + if isinstance(rendered, dict): + render_map = rendered + else: + render_map = {item.page_index: item for item in rendered} + tag_map = {t.page_index: t for t in tags} + label_map: dict[int, str] = {} + if page_labels: + for lbl in page_labels: + label_map[lbl.page] = lbl.kind + + final_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ), + toc_pages, + ) + raw_text_by_page: dict[int, str] = {} + image_path_by_page: dict[int, str] = {} + for page in final_pages: + rend = render_map.get(page) + raw_text_by_page[page] = (rend.raw_text if rend else page_texts.get(page, "")) or "" + if rend and rend.image_path and os.path.exists(rend.image_path): + image_path_by_page[page] = rend.image_path + + rows = build_node_rows( + 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, + verdict=verdict, + vlm_model=vlm_model, + page_assets_by_page=assets_by_page, + node_summary_max_pages=page_memory_config.node_summary_max_pages, + node_assembly_concurrency=page_memory_config.node_assembly_concurrency, + body_text_by_page=body_text_by_page, + ) + return rows, final_pages + + def _build_page_dataframe( *, pdf_path: str, @@ -213,10 +359,10 @@ def _build_page_dataframe( page_memory_config: PageMemoryConfig, trace_recorder: Any | None = None, ) -> pd.DataFrame: - """Build per-page DataFrame via the full C1-C7 pipeline. + """Build per-page DataFrame via Phase A scopes then the full C1-C7 path. Steps: - C4 skeleton_extractor → SectionSkeleton[] + Phase A build_phase_a → frozen scopes / next titles / asset budgets C1 page_renderer → PageRenderResult[] C2 page_plan → PagePlan[] C3 page_tagger → PageTagResult[] @@ -225,91 +371,23 @@ def _build_page_dataframe( C5 page_assets → assets anchored to refined hierarchy pages C7 assemble node-granularity DataFrame """ - from app.services.page_memory.skeleton_extractor import ( - SectionSkeleton, - collapse_single_child_chains, - extract_section_skeletons, - ) - from shared.services.chunks.path_segments import join_document_path - anatomy = getattr(profile, "anatomy", None) + from app.services.page_memory.skeleton_extractor import collapse_single_child_chains + page_count = max(int(profile.page_count or 0), 0) if page_count <= 0: return pd.DataFrame(columns=pd.Index(PARSER_DF_COLUMNS)) asset_extraction_enabled = page_memory_config.asset_extraction_enabled - - # ── C4: skeleton (from profile anatomy) ─────────────────────────── - page_texts = coerce_page_text_cache(getattr(profile, "page_full_text_cache", None)) - with stage_timer("page_memory.skeleton", page_count=page_count): - if anatomy is not None: - skeletons = extract_section_skeletons( - anatomy=anatomy, - filename=filename, - page_texts=page_texts, - ) - else: - skeletons = [] - logger.info( - "[page_memory] C4 skeleton: {} sections from anatomy", - len(skeletons), - ) - - if not skeletons: - skeletons = [ - SectionSkeleton( - section_path=join_document_path([filename, "Root"]), - level=1, - start_page=1, - end_page=page_count, - title="Root", - parent_path=filename, - evidence={"source": "fallback_root", "reason": "no_hierarchy"}, - ) - ] - - 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) - coarse_scopes = _build_hierarchy_scopes( - skeletons=skeletons, + phase_a = build_phase_a( + profile=profile, filename=filename, - page_count=page_count, - ) - toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None - toc_pages = derive_excluded_toc_pages( - skeletons=skeletons, - 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, - "C4.skeleton", - page_info={ - "document_page_count": page_count, - "coarse_scope": page_scope_info(coarse_pages_scope), - }, - variables={ - "sections": _summarize_skeletons(skeletons), - "scope_count": len(coarse_scopes), - "scopes": [ - _scope_manifest( - scope_id=scope.scope_id, - skeletons=scope.skeletons, - page_count=page_count, - strategy=scope.strategy, - ) - for scope in coarse_scopes - ], - }, + trace_recorder=trace_recorder, ) - + page_texts = phase_a.page_texts + coarse_scopes = phase_a.coarse_scopes + next_title_by_path = phase_a.next_title_by_path + toc_pages = phase_a.toc_pages + anatomy = getattr(profile, "anatomy", None) page_features = anatomy.page_features if anatomy else [] page_labels = anatomy.page_labels if anatomy else [] vlm_model = os.environ.get("IMAGE_MODEL") @@ -326,13 +404,6 @@ def _build_page_dataframe( sorted({int(page) for page in toc_pages}), ) scope_results: list[_ScopeRunResult] = [] - asset_pages_remaining = ( - _resolve_asset_max_pages(page_count, page_memory_config) - if asset_extraction_enabled - else 0 - ) - - # Common keyword arguments shared by every scope invocation. _scope_kwargs = dict( pdf_path=pdf_path, filename=filename, @@ -350,39 +421,28 @@ def _build_page_dataframe( ) if scope_concurrency <= 1 or len(coarse_scopes) <= 1: - # ── serial path (identical to previous behaviour) ───────────── for index, scope in enumerate(coarse_scopes, start=1): - result = _run_scope_with_retry( - scope=scope, - scope_index=index, - scope_count=len(coarse_scopes), - asset_max_pages=asset_pages_remaining, - **_scope_kwargs, - ) - scope_results.append(result) - if asset_extraction_enabled: - asset_pages_remaining = max( - asset_pages_remaining - len(result.final_pages), - 0, + scope_results.append( + run_single_scope( + scope=scope, + scope_index=index, + scope_count=len(coarse_scopes), + **_scope_kwargs, ) + ) else: - # ── concurrent path (GeventPool) ────────────────────────────── import gevent from gevent.pool import Pool as GeventPool - asset_allocations = _allocate_asset_pages( - coarse_scopes, asset_pages_remaining, - ) pool = GeventPool( size=min(scope_concurrency, len(coarse_scopes)), ) greenlets = [ pool.spawn( - _run_scope_with_retry, + run_single_scope, scope=scope, scope_index=idx, scope_count=len(coarse_scopes), - asset_max_pages=asset_allocations[idx - 1], **_scope_kwargs, ) for idx, scope in enumerate(coarse_scopes, start=1) @@ -415,48 +475,22 @@ def _build_page_dataframe( ) # ── C7: assemble DataFrame rows ────────────────────────────────── - tag_map = {t.page_index: t for t in tags} - render_map = rendered_map - - # Build page → PageLabel.kind lookup - label_map: dict[int, str] = {} - if page_labels: - for lbl in page_labels: - label_map[lbl.page] = lbl.kind - - # 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 = pages_excluding_toc( - _derive_hierarchy_page_scope( + with stage_timer("page_memory.node_assembly", page_count=page_count): + rows, final_pages_scope = assemble_c7_node_rows( skeletons=skeletons, + rendered=rendered_map, + tags=tags, + assets_by_page=page_assets_by_page, + page_texts=page_texts, + body_text_by_page=getattr(profile, "page_body_text_cache", None), + page_labels=page_labels, + toc_pages=toc_pages, page_count=page_count, - ), - toc_pages, - ) - for page in final_pages_scope: - rend = render_map.get(page) - raw_text_by_page[page] = (rend.raw_text if rend else page_texts.get(page, "")) or "" - if rend and rend.image_path and os.path.exists(rend.image_path): - image_path_by_page[page] = rend.image_path - - from app.services.page_memory.node_assembler import build_node_rows - - with stage_timer("page_memory.node_assembly", page_count=len(final_pages_scope)): - rows = build_node_rows( - 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, verdict=verdict, vlm_model=vlm_model, - page_assets_by_page=page_assets_by_page, - node_summary_max_pages=page_memory_config.node_summary_max_pages, - node_assembly_concurrency=page_memory_config.node_assembly_concurrency, - body_text_by_page=getattr(profile, "page_body_text_cache", None), + page_memory_config=page_memory_config, ) logger.info( "[page_memory] C7 assembled {} node rows (verdict={})", @@ -475,40 +509,12 @@ def _build_page_dataframe( return pd.DataFrame(rows, columns=pd.Index(PARSER_DF_COLUMNS)) -def _build_hierarchy_scopes( - *, - skeletons: list[Any], - filename: str, - page_count: int, -) -> list[CoarseScope]: - return build_hierarchy_scopes( - skeletons=skeletons, - filename=filename, - page_count=page_count, - ) - - -def _allocate_asset_pages( - scopes: list[CoarseScope], total_budget: int, -) -> list[int]: - """Pre-allocate asset page budget proportionally to avoid concurrency races.""" - if total_budget <= 0 or not scopes: - return [0] * len(scopes) - page_counts = [max(s.end_page - s.start_page + 1, 1) for s in scopes] - total_pages = sum(page_counts) - allocations: list[int] = [] - remaining = total_budget - for i, count in enumerate(page_counts): - if i == len(page_counts) - 1: - allocations.append(remaining) - else: - share = min(int(total_budget * count / total_pages), remaining) - allocations.append(share) - remaining -= share - return allocations +_SCOPE_RETRY_DELAY_SECONDS = 10.0 -_SCOPE_RETRY_DELAY_SECONDS = 10.0 +def run_single_scope(**kwargs: Any) -> _ScopeRunResult: + """Run one frozen scope through the existing Stage-3/4 path.""" + return _run_scope_with_retry(**kwargs) def _run_scope_with_retry( @@ -557,7 +563,6 @@ def _run_hierarchy_scope( page_labels: list[Any], vlm_model: str | None, asset_extraction_enabled: bool, - asset_max_pages: int, trace_recorder: Any | None, page_memory_config: PageMemoryConfig, next_title_by_path: dict[str, str | None] | None = None, @@ -578,9 +583,7 @@ def _run_hierarchy_scope( ) scope_skeletons = sort_skeletons(scope.skeletons) - toc_pages = derive_excluded_toc_pages( - skeletons=scope_skeletons, toc_pages=toc_pages, - ) + toc_pages = sorted({int(page) for page in (toc_pages or [])}) scope_manifest = _scope_manifest( scope_id=scope.scope_id, skeletons=scope_skeletons, @@ -761,17 +764,16 @@ def _run_hierarchy_scope( ) assets_by_page: dict[int, list[Any]] = {} - if asset_extraction_enabled and asset_max_pages > 0: + if asset_extraction_enabled: asset_rendered = _select_rendered_pages_with_assets( rendered, page_features ) - with stage_timer("page_memory.assets", page_count=asset_max_pages): + with stage_timer("page_memory.assets", page_count=len(asset_rendered)): assets_by_page = extract_page_assets_from_renders( pdf_path=pdf_path, rendered_pages=asset_rendered, output_dir=output_dir, model_name=page_memory_config.asset_model or resolve_asset_model(), - max_pages=asset_max_pages, confidence_threshold=page_memory_config.asset_confidence_threshold, summary_enabled=page_memory_config.asset_summary_enabled, summary_concurrency=page_memory_config.asset_summary_concurrency, @@ -914,7 +916,7 @@ def _cleanup_page_memory_artifacts(output_dir: str) -> None: path.unlink() except Exception: logger.debug("[page_memory] failed to cleanup artifact {}", path) - for name in ("asset_annotate", "debug", "images", "pages", "scopes", "tables"): + for name in ("asset_annotate", "debug", "images", "pages", "phase_b", "scopes", "tables"): path = root / name try: if path.is_dir(): diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index 43f05d962..fe6896cda 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -74,10 +74,6 @@ def get_asset_model() -> str | None: return resolve_asset_model() or None -def get_asset_max_pages(page_count: int) -> int: - return max(0, page_count) - - def detect_page_assets( *, page: PageRenderResult, @@ -369,14 +365,13 @@ def extract_page_assets_from_renders( rendered_pages: list[PageRenderResult], output_dir: str, model_name: str | None, - max_pages: int, confidence_threshold: float, summary_enabled: bool = False, summary_concurrency: int = 4, table_engine: str = "tabula", table_merge_enabled: bool = True, ) -> dict[int, list[PageAsset]]: - pages = sorted(rendered_pages, key=lambda item: item.page_index)[:max_pages] + pages = sorted(rendered_pages, key=lambda item: item.page_index) source_name = Path(pdf_path).name assets_by_page: dict[int, list[PageAsset]] = {} pending_summaries: list[PageAsset] = [] @@ -906,7 +901,6 @@ def _merge_table_html_files( "build_asset_rows", "extract_page_assets_from_renders", "get_asset_confidence_threshold", - "get_asset_max_pages", "get_asset_model", "merge_cross_page_tables", "page_asset_extraction_enabled", diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index cd2e0ac32..1cac87064 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -46,7 +46,7 @@ class PageTagResult: _MAX_JSON_RETRIES = 1 -_DEFAULT_FINE_MIN_PAGES = 4 +_DEFAULT_FINE_MIN_PAGES = 5 # Dense section-start pages can emit long title JSON; escalate only on truncation. _TITLE_TOKEN_BUDGETS: tuple[int, ...] = (300, 600, 1200) diff --git a/apps/worker/app/services/page_memory/phase_a.py b/apps/worker/app/services/page_memory/phase_a.py new file mode 100644 index 000000000..1db5299b1 --- /dev/null +++ b/apps/worker/app/services/page_memory/phase_a.py @@ -0,0 +1,551 @@ +"""Phase A: coarse skeletons, frozen scope manifest, and coarse page rows.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pandas as pd + +from app.services.document_agent.pdf_text import coerce_page_text_cache +from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import PARSER_DF_COLUMNS +from app.services.page_memory._serialization import ( + derive_hierarchy_page_scope, + scope_manifest, + serialize_skeletons, + write_json, +) +from app.services.page_memory._utils import ( + CoarseScope, + build_hierarchy_scopes, + page_scope_info, +) +from app.services.page_memory.body_text_preview import build_body_text_preview +from app.services.page_memory.fine_hierarchy import build_next_title_by_path +from app.services.page_memory.skeleton_extractor import SectionSkeleton +from shared.services.chunks.path_segments import join_document_path + + +PHASE_B_DIR = "phase_b" +PHASE_B_BLACKBOARD_REL = f"{PHASE_B_DIR}/blackboard.json" +PHASE_B_MANIFEST_REL = f"{PHASE_B_DIR}/manifest.json" +PHASE_B_ANATOMY_REL = f"{PHASE_B_DIR}/anatomy.json" +PHASE_B_FULL_TEXT_REL = f"{PHASE_B_DIR}/page_full_text_cache.json" +PHASE_B_BODY_TEXT_REL = f"{PHASE_B_DIR}/page_body_text_cache.json" +PHASE_B_SOURCE_PDF_REL = "source.pdf" +PHASE_B_INPUT_RELS = ( + PHASE_B_SOURCE_PDF_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_ANATOMY_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, +) +TEXT_ENRICHMENT_STATE = "text" +COARSE_ENRICHMENT_STATE = "coarse" + + +@dataclass(frozen=True) +class PhaseAResult: + coarse_scopes: list[CoarseScope] + next_title_by_path: dict[str, str | None] + toc_pages: list[int] + skeletons: list[SectionSkeleton] + coarse_df: pd.DataFrame + manifest: dict[str, Any] + page_texts: dict[int, str] + + +def build_phase_a( + *, + profile: Any, + filename: str, + trace_recorder: Any | None = None, +) -> PhaseAResult: + """Run through C4 skeletons, freeze scopes, and build coarse page rows.""" + from app.services.page_memory.skeleton_extractor import extract_section_skeletons + + page_texts = coerce_page_text_cache(getattr(profile, "page_full_text_cache", None)) + page_count = max(int(getattr(profile, "page_count", 0) or 0), 0) + anatomy = getattr(profile, "anatomy", None) + if page_count <= 0: + empty = pd.DataFrame(columns=pd.Index(PARSER_DF_COLUMNS)) + return PhaseAResult( + coarse_scopes=[], + next_title_by_path={}, + toc_pages=[], + skeletons=[], + coarse_df=empty, + manifest=_empty_manifest(filename=filename, page_count=page_count), + page_texts=page_texts, + ) + + skeletons = ( + extract_section_skeletons( + anatomy=anatomy, + filename=filename, + page_texts=page_texts, + ) + if anatomy is not None + else [] + ) + if not skeletons: + skeletons = [ + SectionSkeleton( + section_path=join_document_path([filename, "Root"]), + level=1, + start_page=1, + end_page=page_count, + title="Root", + parent_path=filename, + evidence={"source": "fallback_root", "reason": "no_hierarchy"}, + ) + ] + + next_title_by_path = build_next_title_by_path(skeletons) + coarse_scopes = build_hierarchy_scopes( + skeletons=skeletons, + filename=filename, + page_count=page_count, + ) + toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None + toc_pages = sorted( + {int(page) for page in (getattr(toc_result, "toc_pages", None) or [])} + ) + manifest = build_phase_a_manifest( + filename=filename, + page_count=page_count, + skeletons=skeletons, + coarse_scopes=coarse_scopes, + next_title_by_path=next_title_by_path, + toc_pages=toc_pages, + ) + _record_trace_stage( + trace_recorder, + "C4.skeleton", + page_info={ + "document_page_count": page_count, + "coarse_scope": page_scope_info( + pages_excluding_toc( + derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ), + toc_pages, + ) + ), + }, + variables={ + "sections": [ + { + "path": item.section_path, + "level": item.level, + "start_page": item.start_page, + "end_page": item.end_page, + } + for item in skeletons + ], + "scope_count": len(coarse_scopes), + "scopes": [entry for entry in manifest["scopes"]], + }, + ) + coarse_df = build_coarse_dataframe( + coarse_scopes=coarse_scopes, + page_texts=page_texts, + body_text_by_page=coerce_page_text_cache( + getattr(profile, "page_body_text_cache", None) + ), + toc_pages=toc_pages, + page_count=page_count, + filename=filename, + ) + return PhaseAResult( + coarse_scopes=coarse_scopes, + next_title_by_path=next_title_by_path, + toc_pages=toc_pages, + skeletons=skeletons, + coarse_df=coarse_df, + manifest=manifest, + page_texts=page_texts, + ) + + +def build_phase_a_manifest( + *, + filename: str, + page_count: int, + skeletons: list[SectionSkeleton], + coarse_scopes: list[CoarseScope], + next_title_by_path: dict[str, str | None], + toc_pages: list[int], +) -> dict[str, Any]: + scope_entries: list[dict[str, Any]] = [] + for scope in coarse_scopes: + base = scope_manifest( + scope_id=scope.scope_id, + skeletons=scope.skeletons, + page_count=page_count, + strategy=scope.strategy, + ) + skeleton_paths = [ + str(getattr(item, "section_path", "") or "") for item in scope.skeletons + ] + processing_pages = pages_excluding_toc( + derive_hierarchy_page_scope( + skeletons=scope.skeletons, + page_count=page_count, + ), + toc_pages, + ) + excluded_toc_pages = sorted( + page + for page in toc_pages + if scope.start_page <= page <= scope.end_page + ) + scope_entries.append( + { + **base, + "start_page": scope.start_page, + "end_page": scope.end_page, + "skeletons": serialize_skeletons(scope.skeletons), + "next_title_slice": { + path: next_title_by_path.get(path) + for path in skeleton_paths + if path in next_title_by_path + }, + "toc_pages": list(toc_pages), + "processing_pages": processing_pages, + "excluded_toc_pages": excluded_toc_pages, + } + ) + return { + "filename": filename, + "page_count": page_count, + "toc_pages": list(toc_pages), + "next_title_by_path": dict(next_title_by_path), + "skeletons": serialize_skeletons(skeletons), + "scopes": scope_entries, + } + + +def build_coarse_dataframe( + *, + coarse_scopes: list[CoarseScope], + page_texts: dict[int, str], + body_text_by_page: dict[int, str], + toc_pages: list[int], + page_count: int, + filename: str, +) -> pd.DataFrame: + """One coarse page row per scoped skeleton: raw cache text, no VLM.""" + rows: list[dict[str, Any]] = [] + for scope in coarse_scopes: + for skeleton in scope.skeletons: + pages = _skeleton_pages( + skeleton=skeleton, + page_count=page_count, + toc_pages=toc_pages, + ) + content = "\n\n".join(page_texts.get(page, "") for page in pages).strip() + path = str(skeleton.section_path) + rows.append( + { + "content": content, + "path": path, + "type": "page", + "length": len(content), + "keywords": "", + "summary": "", + "know_id": f"node_{gen_str_codes(f'{filename}::{path}')}", + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": ",".join(str(page) for page in pages), + "extra_metadata": { + "scope_id": scope.scope_id, + "enrichment_state": COARSE_ENRICHMENT_STATE, + "body_text_preview": build_body_text_preview( + body_text_by_page.get(page, "") for page in pages + ), + }, + } + ) + return pd.DataFrame(rows, columns=pd.Index(PARSER_DF_COLUMNS)) + + +def build_coarse_text_dataframe( + *, + page_texts: dict[int, str], + body_text_by_page: dict[int, str], + page_count: int, + filename: str, +) -> pd.DataFrame: + """M0 baseline: one flat per-page row (path=filename/p{n}, enrichment_state='text').""" + rows: list[dict[str, Any]] = [] + for page in range(1, max(int(page_count or 0), 0) + 1): + content = str(page_texts.get(page, "") or "").strip() + scope_id = f"p{page}-{page}" + path = join_document_path([filename, f"p{page}"]) + rows.append( + { + "content": content, + "path": path, + "type": "page", + "length": len(content), + "keywords": "", + "summary": "", + "know_id": f"node_{gen_str_codes(f'{filename}::{path}')}", + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": str(page), + "extra_metadata": { + "scope_id": scope_id, + "enrichment_state": TEXT_ENRICHMENT_STATE, + "body_text_preview": build_body_text_preview( + [body_text_by_page.get(page, "")] + ), + }, + } + ) + return pd.DataFrame(rows, columns=pd.Index(PARSER_DF_COLUMNS)) + + +def build_visual_page_dataframe( + *, + page_count: int, + filename: str, + image_path_by_page: dict[int, str], + embeddings_by_page: dict[int, list[float]], + output_dir: str, +) -> pd.DataFrame: + """Pure-visual M0 baseline for scanned PDFs: empty text, page_assets + _page_embedding.""" + from app.services.page_memory.node_assembler import _build_page_citation_assets + + rows: list[dict[str, Any]] = [] + for page in range(1, max(int(page_count or 0), 0) + 1): + scope_id = f"p{page}-{page}" + path = join_document_path([filename, f"p{page}"]) + page_assets = _build_page_citation_assets( + pages=[page], + image_path_by_page=image_path_by_page, + output_dir=output_dir, + ) + extra_metadata: dict[str, Any] = { + "scope_id": scope_id, + "enrichment_state": TEXT_ENRICHMENT_STATE, + "body_text_preview": "", + } + if page_assets: + extra_metadata["page_assets"] = page_assets + embedding = embeddings_by_page.get(page) + if isinstance(embedding, list) and embedding: + extra_metadata["_page_embedding"] = [float(val) for val in embedding] + rows.append( + { + "content": "", + "path": path, + "type": "page", + "length": 0, + "keywords": "", + "summary": "", + "know_id": f"node_{gen_str_codes(f'{filename}::{path}')}", + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": str(page), + "extra_metadata": extra_metadata, + } + ) + return pd.DataFrame(rows, columns=pd.Index(PARSER_DF_COLUMNS)) + + +def persist_m0_inputs( + *, + output_dir: str, + profile: Any, +) -> list[str]: + """Persist M0 blackboard snapshot + text caches under ``output_dir/phase_b/``.""" + from app.services.document_agent.pdf_text import PageTextBands + + root = Path(output_dir) + blackboard = getattr(profile, "blackboard", None) + if blackboard is not None: + doc_profile = getattr(blackboard, "document_profile", None) + bands_raw = getattr(blackboard, "page_full_text_cache", None) or {} + bands_json = { + str(page): PageTextBands.from_any(val).to_dict() + for page, val in dict(bands_raw).items() + } + snapshot = { + "version": "1.0", + "page_count": int(getattr(blackboard, "page_count", 0) or 0), + "document_profile": ( + doc_profile.to_dict() if doc_profile is not None else None + ), + "page_features": [ + feature.to_dict() + for feature in (getattr(blackboard, "page_features", None) or []) + ], + "page_labels": [ + label.to_dict() + for label in (getattr(blackboard, "page_labels", None) or []) + ], + "doc_stats": dict(getattr(blackboard, "doc_stats", None) or {}), + "global_signals": dict(getattr(blackboard, "global_signals", None) or {}), + "page_full_text_bands": bands_json, + } + write_json(root / PHASE_B_BLACKBOARD_REL, snapshot) + write_json( + root / PHASE_B_FULL_TEXT_REL, + _jsonable_page_map(getattr(profile, "page_full_text_cache", None)), + ) + write_json( + root / PHASE_B_BODY_TEXT_REL, + _jsonable_page_map(getattr(profile, "page_body_text_cache", None)), + ) + return existing_phase_b_input_refs(output_dir) + + +def persist_phase_b_inputs( + *, + output_dir: str, + profile: Any, + manifest: dict[str, Any], +) -> list[str]: + """Write Phase B inputs under ``output_dir`` and return existing relative paths.""" + root = Path(output_dir) + write_json(root / PHASE_B_MANIFEST_REL, manifest) + write_json( + root / PHASE_B_FULL_TEXT_REL, + _jsonable_page_map(getattr(profile, "page_full_text_cache", None)), + ) + write_json( + root / PHASE_B_BODY_TEXT_REL, + _jsonable_page_map(getattr(profile, "page_body_text_cache", None)), + ) + anatomy = getattr(profile, "anatomy", None) + if anatomy is not None and hasattr(anatomy, "to_dict"): + write_json(root / PHASE_B_ANATOMY_REL, anatomy.to_dict()) + return existing_phase_b_input_refs(output_dir) + + +def existing_phase_b_input_refs(output_dir: str) -> list[str]: + return [ + relative + for relative in PHASE_B_INPUT_RELS + if os.path.isfile(os.path.join(output_dir, relative)) + ] + + +def scopes_from_phase_a_manifest(manifest: dict[str, Any]) -> list[dict[str, Any]]: + scopes: list[dict[str, Any]] = [] + for entry in manifest.get("scopes") or []: + if not isinstance(entry, dict): + continue + scopes.append( + { + "scope_id": str(entry["scope_id"]), + "skeletons": skeletons_from_rows(list(entry.get("skeletons") or [])), + "start_page": int(entry["start_page"]), + "end_page": int(entry["end_page"]), + "strategy": str(entry.get("strategy") or ""), + "processing_pages": [ + int(page) for page in (entry.get("processing_pages") or []) + ], + "excluded_toc_pages": [ + int(page) for page in (entry.get("excluded_toc_pages") or []) + ], + } + ) + return scopes + + +def skeletons_from_phase_a_manifest(manifest: dict[str, Any]) -> list[SectionSkeleton]: + return skeletons_from_rows(list(manifest.get("skeletons") or [])) + + +def skeletons_from_rows(rows: list[Any]) -> list[SectionSkeleton]: + skeletons: list[SectionSkeleton] = [] + for row in rows: + if not isinstance(row, dict): + continue + skeletons.append( + SectionSkeleton( + section_path=str(row["section_path"]), + title=str(row["title"]), + level=int(row["level"]), + start_page=int(row["start_page"]), + end_page=int(row["end_page"]), + parent_path=row.get("parent_path"), + evidence=dict(row.get("evidence") or {}), + ) + ) + return skeletons + + +def next_title_by_path_from_manifest(manifest: dict[str, Any]) -> dict[str, str | None]: + raw = manifest.get("next_title_by_path") + if not isinstance(raw, dict): + raise ValueError("phase A manifest is missing next_title_by_path") + next_map: dict[str, str | None] = {} + for path, title in raw.items(): + next_map[str(path)] = None if title is None else str(title) + return next_map + + +def require_phase_a_manifest(stage2: dict[str, Any] | None) -> dict[str, Any]: + if not isinstance(stage2, dict): + raise ValueError("pipeline state is missing stages.stage2") + manifest = stage2.get("scope_manifest") + if not isinstance(manifest, dict) or "scopes" not in manifest: + raise ValueError("pipeline state is missing stages.stage2.scope_manifest") + return manifest + + +def _empty_manifest(*, filename: str, page_count: int) -> dict[str, Any]: + return { + "filename": filename, + "page_count": page_count, + "toc_pages": [], + "next_title_by_path": {}, + "skeletons": [], + "scopes": [], + } + + +def _skeleton_pages( + *, + skeleton: SectionSkeleton, + page_count: int, + toc_pages: list[int], +) -> list[int]: + start_page = max(1, int(getattr(skeleton, "start_page", 1) or 1)) + end_page = min( + page_count, + int(getattr(skeleton, "end_page", start_page) or start_page), + ) + if end_page < start_page or page_count <= 0: + return [] + return pages_excluding_toc(list(range(start_page, end_page + 1)), toc_pages) + + +def _jsonable_page_map(raw: Any) -> dict[str, str]: + coerced = coerce_page_text_cache(raw) + return {str(page): text for page, text in sorted(coerced.items())} + + +def _record_trace_stage( + trace_recorder: Any | None, + stage: str, + *, + page_info: dict[str, Any], + variables: dict[str, Any], +) -> None: + if trace_recorder is None or not hasattr(trace_recorder, "record_stage"): + return + trace_recorder.record_stage(stage, page_info=page_info, variables=variables) diff --git a/apps/worker/app/services/page_memory/scope_enrichment_service.py b/apps/worker/app/services/page_memory/scope_enrichment_service.py new file mode 100644 index 000000000..f1ae2d1a4 --- /dev/null +++ b/apps/worker/app/services/page_memory/scope_enrichment_service.py @@ -0,0 +1,1248 @@ +"""Phase B per-scope enrichment service for page-memory documents.""" + +from __future__ import annotations + +import json +import os +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, cast + +import pandas as pd +from loguru import logger +from sqlalchemy import func, select, update + +from app.services.connect_builder.summary_builder import ( + build_section_summary_lookup, + enrich_doc_nav_summaries, + ensure_doc_nav_json, +) +from app.services.document_agent.manifest import PageFeature, PageLabel +from app.services.document_agent.pdf_text import coerce_page_text_cache +from app.services.document_ingestion.artifact_refs import ( + collect_referenced_artifact_refs, +) +from app.services.document_parser.support.parser_rows import PARSER_DF_COLUMNS +from app.services.page_memory._utils import CoarseScope, sort_skeletons +from app.services.page_memory.memory_service import ( + assemble_c7_node_rows, + run_single_scope, +) +from app.services.page_memory.phase_a import ( + PHASE_B_ANATOMY_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_BODY_TEXT_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_SOURCE_PDF_REL, + build_phase_a, + next_title_by_path_from_manifest, + persist_phase_b_inputs, + skeletons_from_rows, +) +from app.services.page_memory.skeleton_extractor import ( + SectionSkeleton, + collapse_single_child_chains, +) +from shared.core.database_sync import get_sync_db_context +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentScopeEnrichment, +) +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.models.schemas.page_memory_config import PageMemoryConfig +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.ai.llm_overrides import cleanup_llm_overrides, init_llm_overrides +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.services.jobs.lifecycle.scope_enrichments import ( + enqueue_low_priority_scope_enrichments_after_commit, + scope_claimable_filter, + write_coarse_scope_enrichments, +) +from shared.services.redis.redis_sync_service import ( + SyncJobMetadataService, + SyncRedisServiceFactory, +) +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.storage.result_storage import ResultStorage, get_result_storage +from shared.services.storage.zip_result_service import ZipResultService +from shared.utils.utc_now import utc_now_naive + +REFINED_ENRICHMENT_STATE = "refined" + +ScopePublisher = Callable[..., Any] + + +@dataclass(frozen=True) +class _ClaimedScopeTarget: + document_id: str + job_id: str + user_id: str + namespace: str + source_file_name: str + raw_prefix: str | None + job_metadata: dict[str, Any] + page_memory_config: PageMemoryConfig + + +@dataclass(frozen=True) +class ScopeRefinementOutput: + scope_id: str + dataframe: pd.DataFrame + chunks: list[dict[str, Any]] + section_summaries: dict[str, str] + artifact_refs: set[str] + + +def execute_structure_enrichment( + *, + job_result_id: str, + result_storage_factory: Callable[[], ResultStorage] = get_result_storage, + scope_publisher: ScopePublisher | None = None, +) -> dict[str, Any]: + """Resume profiling (TOC + anatomy + Phase A), restructure-publish coarse, and enqueue fine scopes.""" + claim_status, target = _load_and_validate_structure_target( + job_result_id=job_result_id, + ) + if target is None: + return { + "status": claim_status, + "job_result_id": job_result_id, + } + + storage = result_storage_factory() + publisher = scope_publisher or _resolve_default_scope_publisher() + workspace_dir = tempfile.mkdtemp(prefix=f"pm_struct_{job_result_id}_") + init_llm_overrides(JobMetadataHelper.get_llm_config(target.job_metadata)) + try: + add_dir = os.path.join( + workspace_dir, + Path(target.source_file_name).stem or "document", + ) + os.makedirs(add_dir, exist_ok=True) + _download_phase_b_inputs( + storage=storage, + job_id=target.job_id, + raw_prefix=target.raw_prefix, + job_metadata=target.job_metadata, + output_dir=add_dir, + ) + profile = _resume_profile_from_m0_workspace( + add_dir=add_dir, + target=target, + ) + phase_a = build_phase_a( + profile=profile, + filename=target.source_file_name, + trace_recorder=None, + ) + phase_b_refs = persist_phase_b_inputs( + output_dir=add_dir, + profile=profile, + manifest=phase_a.manifest, + ) + _upload_scope_artifacts( + storage=storage, + job_id=target.job_id, + raw_prefix=target.raw_prefix, + workspace_dir=workspace_dir, + artifact_refs={ + PHASE_B_MANIFEST_REL, + PHASE_B_ANATOMY_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, + } + & set(phase_b_refs), + ) + coarse_chunks = cast( + list[dict[str, Any]], dataframe_to_chunks(phase_a.coarse_df) + ) + coarse_summaries = _build_scope_section_summaries( + output_dir=add_dir, + chunks=coarse_chunks, + source_file_name=target.source_file_name, + summary_use_llm=bool( + JobMetadataHelper.get_parsing_param( + target.job_metadata, "summary_use_llm", False + ) + ), + ) + published_ok = _publish_coarse_structure_and_write_scopes( + job_result_id=job_result_id, + target=target, + chunks=coarse_chunks, + section_summaries=coarse_summaries, + phase_b_refs=phase_b_refs, + scope_publisher=publisher, + ) + if not published_ok: + return {"status": "aborted", "job_result_id": job_result_id} + + enqueued_scopes = enqueue_low_priority_scope_enrichments_after_commit( + job_result_id + ) + return { + "status": "coarse", + "job_result_id": job_result_id, + "scope_count": len(phase_a.coarse_scopes), + "enqueued_scopes": enqueued_scopes, + "chunk_count": len(coarse_chunks), + } + except Exception: + _release_structure_claim(job_id=target.job_id, job_result_id=job_result_id) + raise + finally: + cleanup_llm_overrides() + shutil.rmtree(workspace_dir, ignore_errors=True) + + +def execute_scope_enrichment( + *, + job_result_id: str, + scope_id: str, + result_storage_factory: Callable[[], ResultStorage] = get_result_storage, + scope_publisher: ScopePublisher | None = None, +) -> dict[str, Any]: + """Claim one scope, run C1-C7 on its subtree, and publish its refined rows.""" + claim_status, target = _claim_and_validate_scope( + job_result_id=job_result_id, + scope_id=scope_id, + ) + if target is None: + return { + "status": claim_status, + "job_result_id": job_result_id, + "scope_id": scope_id, + } + + storage = result_storage_factory() + publisher = scope_publisher or _resolve_default_scope_publisher() + workspace_dir = tempfile.mkdtemp(prefix=f"pm_scope_{scope_id}_") + init_llm_overrides(JobMetadataHelper.get_llm_config(target.job_metadata)) + try: + output = _run_scope_in_workspace( + workspace_dir=workspace_dir, + target=target, + scope_id=scope_id, + storage=storage, + ) + _upload_scope_artifacts( + storage=storage, + job_id=target.job_id, + raw_prefix=target.raw_prefix, + workspace_dir=workspace_dir, + artifact_refs=output.artifact_refs, + ) + all_refined = _publish_and_mark_refined( + job_result_id=job_result_id, + scope_id=scope_id, + target=target, + output=output, + scope_publisher=publisher, + ) + if all_refined is None: + return { + "status": "aborted", + "job_result_id": job_result_id, + "scope_id": scope_id, + } + if all_refined: + _finalize_all_scopes_refined( + job_result_id=job_result_id, + target=target, + storage=storage, + workspace_dir=workspace_dir, + ) + return { + "status": REFINED_ENRICHMENT_STATE, + "job_result_id": job_result_id, + "scope_id": scope_id, + "chunk_count": len(output.chunks), + "all_refined": all_refined, + } + except Exception: + _mark_scope_failed(job_result_id=job_result_id, scope_id=scope_id) + raise + finally: + cleanup_llm_overrides() + shutil.rmtree(workspace_dir, ignore_errors=True) + + +def run_single_scope_refinement( + *, + scope_id: str, + pdf_path: str, + output_dir: str, + manifest: dict[str, Any], + page_texts: dict[int, str], + body_text_by_page: dict[int, str], + page_features: list[PageFeature], + page_labels: list[PageLabel], + page_memory_config: PageMemoryConfig, + source_file_name: str | None = None, + summary_use_llm: bool = False, +) -> ScopeRefinementOutput: + """Run C1-C7 for a single frozen scope and convert its subtree to chunks.""" + entry, scope_index, scope_count = _find_manifest_scope(manifest, scope_id) + filename = str(source_file_name or manifest["filename"]) + page_count = max(int(manifest.get("page_count") or 0), 0) + toc_pages = [int(page) for page in (manifest.get("toc_pages") or [])] + next_title_by_path = next_title_by_path_from_manifest(manifest) + + coarse_scope = CoarseScope( + scope_id=scope_id, + skeletons=skeletons_from_rows(list(entry.get("skeletons") or [])), + strategy=str(entry["strategy"]), + start_page=int(entry["start_page"]), + end_page=int(entry["end_page"]), + ) + vlm_model = os.environ.get("IMAGE_MODEL") + + scope_result = run_single_scope( + scope=coarse_scope, + scope_index=scope_index, + scope_count=scope_count, + pdf_path=pdf_path, + filename=filename, + output_dir=output_dir, + page_count=page_count, + page_texts=page_texts, + page_features=page_features, + page_labels=page_labels, + vlm_model=vlm_model, + asset_extraction_enabled=page_memory_config.asset_extraction_enabled, + trace_recorder=None, + page_memory_config=page_memory_config, + next_title_by_path=next_title_by_path, + toc_pages=toc_pages, + ) + + collapsed_skeletons = collapse_single_child_chains( + sort_skeletons(list(scope_result.hierarchy)) + ) + rows = _assemble_scope_rows( + scope_id=scope_id, + skeletons=collapsed_skeletons, + rendered=scope_result.rendered, + tags=scope_result.tags, + assets_by_page=scope_result.assets_by_page, + page_texts=page_texts, + body_text_by_page=body_text_by_page, + page_labels=page_labels, + toc_pages=toc_pages, + page_count=page_count, + output_dir=output_dir, + filename=filename, + vlm_model=vlm_model, + page_memory_config=page_memory_config, + ) + dataframe = pd.DataFrame(rows, columns=pd.Index(PARSER_DF_COLUMNS)) + chunks = cast(list[dict[str, Any]], dataframe_to_chunks(dataframe)) + section_summaries = _build_scope_section_summaries( + output_dir=output_dir, + chunks=chunks, + source_file_name=filename, + summary_use_llm=summary_use_llm, + ) + artifact_refs = collect_referenced_artifact_refs(chunks) + return ScopeRefinementOutput( + scope_id=scope_id, + dataframe=dataframe, + chunks=chunks, + section_summaries=section_summaries, + artifact_refs=artifact_refs, + ) + + +def _assemble_scope_rows( + *, + scope_id: str, + skeletons: list[SectionSkeleton], + rendered: list[Any], + tags: list[Any], + assets_by_page: dict[int, list[Any]], + page_texts: dict[int, str], + body_text_by_page: dict[int, str], + page_labels: list[PageLabel], + toc_pages: list[int], + page_count: int, + output_dir: str, + filename: str, + vlm_model: str | None, + page_memory_config: PageMemoryConfig, +) -> list[dict[str, Any]]: + rows, _final_pages = assemble_c7_node_rows( + skeletons=skeletons, + rendered=rendered, + tags=tags, + assets_by_page=assets_by_page, + page_texts=page_texts, + body_text_by_page=body_text_by_page, + page_labels=page_labels, + toc_pages=toc_pages, + page_count=page_count, + output_dir=output_dir, + filename=filename, + verdict="page", + vlm_model=vlm_model, + page_memory_config=page_memory_config, + ) + for row in rows: + extra = row.get("extra_metadata") + extra_dict = dict(extra) if isinstance(extra, dict) else {} + extra_dict["scope_id"] = scope_id + if str(row.get("type") or "") == "page": + extra_dict["enrichment_state"] = REFINED_ENRICHMENT_STATE + row["extra_metadata"] = extra_dict + return rows + + +def _build_claimed_target( + *, + document: Document, + job: Job, + job_metadata: dict[str, Any], + source_file_name: str, +) -> _ClaimedScopeTarget: + raw_prefix_val = job_metadata.get("result_raw_prefix") + raw_prefix = ( + str(raw_prefix_val) + if isinstance(raw_prefix_val, str) and raw_prefix_val + else None + ) + corpus_storage = CorpusStorage.resolve_namespace(document.namespace) + return _ClaimedScopeTarget( + document_id=str(document.document_id), + job_id=str(job.job_id), + user_id=corpus_storage.resolve_owner(str(job.user_id)), + namespace=normalize_retrieval_namespace(document.namespace), + source_file_name=source_file_name, + raw_prefix=raw_prefix, + job_metadata=job_metadata, + page_memory_config=JobMetadataHelper.get_page_memory_config(job_metadata), + ) + + +def _structure_claim_expired( + *, + job_result: JobResult, + job: Job, + now: Any, +) -> bool: + doc_meta = job_result.document_metadata if isinstance(job_result.document_metadata, dict) else {} + job_meta = job.job_metadata if isinstance(job.job_metadata, dict) else {} + raw_claimed_at = doc_meta.get("structure_claimed_at") or job_meta.get("structure_claimed_at") + if not raw_claimed_at: + return True + try: + from datetime import datetime, timedelta + from shared.core.config import app_config + + claimed_at = datetime.fromisoformat(str(raw_claimed_at).replace("Z", "+00:00")).replace(tzinfo=None) + return (now - claimed_at) >= timedelta(seconds=app_config.TASK_TIME_LIMIT_SECONDS) + except Exception: + return True + + +def _release_structure_claim(*, job_id: str, job_result_id: str) -> None: + with get_sync_db_context() as db: + job = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()).scalar_one_or_none() + if job is not None and isinstance(job.job_metadata, dict): + meta = dict(job.job_metadata) + if str(meta.get("enrichment_status") or "").strip().lower() == "text": + meta.pop("structure_claimed_at", None) + job.job_metadata = meta + jr = ( + db.execute( + select(JobResult) + .where(JobResult.id == job_result_id) + .with_for_update(of=JobResult) + ).scalar_one_or_none() + ) + if jr is not None and isinstance(jr.document_metadata, dict): + doc_meta = dict(jr.document_metadata) + if str(doc_meta.get("enrichment_status") or "").strip().lower() == "text": + doc_meta.pop("structure_claimed_at", None) + jr.document_metadata = doc_meta + + +def _load_and_validate_structure_target( + *, + job_result_id: str, +) -> tuple[str, _ClaimedScopeTarget | None]: + now = utc_now_naive() + with get_sync_db_context() as db: + job_result = db.execute( + select(JobResult) + .where(JobResult.id == job_result_id) + .with_for_update(of=JobResult) + ).scalar_one_or_none() + if job_result is None: + return "aborted", None + doc_id = job_result.document_id or job_result.demo_document_id + document = ( + db.execute( + select(Document).where(Document.document_id == str(doc_id)) + ).scalar_one_or_none() + if doc_id + else None + ) + job = db.execute( + select(Job).where(Job.job_id == job_result.job_id).with_for_update() + ).scalar_one_or_none() + job_metadata = dict(job.job_metadata or {}) if job is not None else {} + source_file_name = ( + ( + document.source_file_name + or JobMetadataHelper.get_source_file_name(job_metadata) + ) + if document is not None + else None + ) + if ( + document is None + or document.status != "active" + or document.current_job_result_id != job_result_id + or job is None + or not source_file_name + ): + return "aborted", None + enrichment_status = ( + str( + (job_result.document_metadata or {}).get("enrichment_status") + or job_metadata.get("enrichment_status") + or "" + ) + .strip() + .lower() + ) + if enrichment_status != "text": + return "skipped", None + if not _structure_claim_expired(job_result=job_result, job=job, now=now): + return "skipped", None + claimed_iso = now.isoformat() + job_metadata["structure_claimed_at"] = claimed_iso + job.job_metadata = job_metadata + job.updated_at = now + job_result.document_metadata = { + **dict(job_result.document_metadata or {}), + "structure_claimed_at": claimed_iso, + } + return ( + "claimed", + _build_claimed_target( + document=document, + job=job, + job_metadata=job_metadata, + source_file_name=source_file_name, + ), + ) + + +def _resume_profile_from_m0_workspace( + *, + add_dir: str, + target: _ClaimedScopeTarget, +) -> Any: + from app.services.document_agent.coordinator import ProfileCoordinator + from app.services.document_agent.manifest import DocumentProfile + from app.services.document_agent.pdf_text import ( + page_bands_map, + page_body_content_map, + page_content_map, + ) + from app.services.document_parser.profiling.doc_profiler import _map_toc_profile + from app.services.document_parser.profiling.domain_taxonomy import OTHER_CODE + from app.services.document_parser.profiling.profile_model import ( + ParserDocumentProfile, + ) + from app.services.document_parser.profiling.taxonomy import PdfRoutingCategory + from shared.core.config import settings + + pdf_path = os.path.join(add_dir, PHASE_B_SOURCE_PDF_REL) + bb_path = os.path.join(add_dir, PHASE_B_BLACKBOARD_REL) + coordinator = ProfileCoordinator( + pdf_path=pdf_path, + job_id=target.job_id, + output_dir=os.path.join(add_dir, "_doc_agent"), + db=None, + model=settings.IMAGE_MODEL, + settings={ + "vlm_model": settings.IMAGE_MODEL, + "toc_profile_enabled": True, + "skip_toc_anchoring": False, + "model": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL, + }, + ) + bb_raw = _read_json_dict(bb_path) + bb = coordinator.blackboard + bb.page_count = int(bb_raw.get("page_count") or 0) + dp_raw = bb_raw.get("document_profile") or {} + doc_profile = DocumentProfile( + is_scanned=bool(dp_raw.get("is_scanned", False)), + category=list(dp_raw.get("category") or [OTHER_CODE]), + category_label=str(dp_raw.get("category_label") or "unknown"), + routing_category=str(dp_raw.get("routing_category") or "generic"), + language=str(dp_raw.get("language") or "unknown"), + header_y=dp_raw.get("header_y"), + footer_y=dp_raw.get("footer_y"), + ) + bb.document_profile = doc_profile + coordinator._coarse_profile_cache = doc_profile + features, labels = _parse_features_and_labels( + bb_raw.get("page_features") or [], + bb_raw.get("page_labels") or [], + ) + bb.page_features = features + bb.page_labels = labels + bb.doc_stats = dict(bb_raw.get("doc_stats") or {}) + bb.global_signals = dict(bb_raw.get("global_signals") or {}) + bb.page_full_text_cache = page_bands_map(bb_raw["page_full_text_bands"]) + bb.page_text_search_view = None + if doc_profile.is_scanned and not bb.page_full_text_cache: + coordinator._run_text_scan() + coordinator._ensure_toc_profile(strict=False) + + routing_category = PdfRoutingCategory.normalize(doc_profile.routing_category) + profile = ParserDocumentProfile( + file_type="pdf", + category=doc_profile.category, + category_label=doc_profile.category_label, + routing_category=routing_category, + is_scanned=doc_profile.is_scanned, + page_count=coordinator.blackboard.page_count, + language=doc_profile.language, + metrics={ + "doc_stats": coordinator.blackboard.doc_stats, + "doc_shape": coordinator.blackboard.global_signals.get("doc_shape", {}), + "page_kind_counts": coordinator.blackboard.global_signals.get( + "page_kind_counts", {} + ), + }, + page_full_text_cache=page_content_map( + coordinator.blackboard.page_full_text_cache + ), + page_body_text_cache=page_body_content_map( + coordinator.blackboard.page_full_text_cache + ), + ) + if not profile.is_atlas: + if profile.page_count > settings.MAX_PDF_PAGE_LIMIT: + profile.anatomy = coordinator.run_structural(skip_shard_plan=True) + else: + profile.anatomy = coordinator.run_lightweight_anatomy( + skip_shard_plan=True + ) + profile.toc = _map_toc_profile(coordinator) + return profile + + +def _publish_coarse_structure_and_write_scopes( + *, + job_result_id: str, + target: _ClaimedScopeTarget, + chunks: list[dict[str, Any]], + section_summaries: dict[str, str], + phase_b_refs: list[str], + scope_publisher: ScopePublisher, +) -> bool: + publication_scope = DocumentPublicationScope( + user_id=target.user_id, + namespace=target.namespace, + document_id=target.document_id, + job_result_id=job_result_id, + source_file_name=target.source_file_name, + ) + with get_sync_db_context() as db: + document = db.execute( + select(Document) + .where(Document.document_id == target.document_id) + .with_for_update() + ).scalar_one_or_none() + if ( + document is None + or document.status != "active" + or document.current_job_result_id != job_result_id + ): + return False + scope_publisher( + db, + scope=publication_scope, + scope_id="*", + chunks=chunks, + section_summaries=section_summaries, + ) + write_coarse_scope_enrichments( + db, + document_id=target.document_id, + job_result_id=job_result_id, + ) + _update_job_and_result_metadata( + job_id=target.job_id, + job_result_id=job_result_id, + updates={ + "enrichment_status": "coarse", + "page_memory_phase_b": phase_b_refs, + }, + ) + return True + + +def _claim_and_validate_scope( + *, + job_result_id: str, + scope_id: str, +) -> tuple[str, _ClaimedScopeTarget | None]: + now = utc_now_naive() + with get_sync_db_context() as db: + claimed_doc_id = db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .where(scope_claimable_filter(now)) + .values( + state="running", + attempts=DocumentScopeEnrichment.attempts + 1, + updated_at=now, + ) + .returning(DocumentScopeEnrichment.document_id) + ).scalar_one_or_none() + if claimed_doc_id is None: + return "skipped", None + + document = db.execute( + select(Document).where(Document.document_id == str(claimed_doc_id)) + ).scalar_one_or_none() + job_result = db.execute( + select(JobResult).where(JobResult.id == job_result_id) + ).scalar_one_or_none() + job = ( + db.execute( + select(Job).where(Job.job_id == job_result.job_id) + ).scalar_one_or_none() + if job_result is not None + else None + ) + job_metadata = dict(job.job_metadata or {}) if job is not None else {} + source_file_name = ( + ( + document.source_file_name + or JobMetadataHelper.get_source_file_name(job_metadata) + ) + if document is not None + else None + ) + if ( + document is None + or document.status != "active" + or document.current_job_result_id != job_result_id + or job_result is None + or job is None + or not source_file_name + ): + db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .values(state="failed", updated_at=now) + ) + return "aborted", None + + return ( + "claimed", + _build_claimed_target( + document=document, + job=job, + job_metadata=job_metadata, + source_file_name=source_file_name, + ), + ) + + +def _run_scope_in_workspace( + *, + workspace_dir: str, + target: _ClaimedScopeTarget, + scope_id: str, + storage: ResultStorage, +) -> ScopeRefinementOutput: + add_dir = os.path.join( + workspace_dir, + Path(target.source_file_name).stem or "document", + ) + os.makedirs(add_dir, exist_ok=True) + _download_phase_b_inputs( + storage=storage, + job_id=target.job_id, + raw_prefix=target.raw_prefix, + job_metadata=target.job_metadata, + output_dir=add_dir, + ) + manifest = _read_json_dict(os.path.join(add_dir, PHASE_B_MANIFEST_REL)) + page_texts = coerce_page_text_cache( + _read_json_dict(os.path.join(add_dir, PHASE_B_FULL_TEXT_REL)) + ) + body_text_by_page = coerce_page_text_cache( + _read_json_dict(os.path.join(add_dir, PHASE_B_BODY_TEXT_REL)) + ) + page_features, page_labels = _load_anatomy_features_and_labels( + os.path.join(add_dir, PHASE_B_ANATOMY_REL) + ) + summary_use_llm = bool( + JobMetadataHelper.get_parsing_param( + target.job_metadata, + "summary_use_llm", + False, + ) + ) + return run_single_scope_refinement( + scope_id=scope_id, + pdf_path=os.path.join(add_dir, PHASE_B_SOURCE_PDF_REL), + output_dir=add_dir, + manifest=manifest, + page_texts=page_texts, + body_text_by_page=body_text_by_page, + page_features=page_features, + page_labels=page_labels, + page_memory_config=target.page_memory_config, + source_file_name=target.source_file_name, + summary_use_llm=summary_use_llm, + ) + + +def _download_phase_b_inputs( + *, + storage: ResultStorage, + job_id: str, + raw_prefix: str | None, + job_metadata: dict[str, Any], + output_dir: str, +) -> None: + configured_refs = job_metadata.get("page_memory_phase_b") + if not isinstance(configured_refs, list) or not configured_refs: + raise ValueError("job metadata is missing the page_memory_phase_b input list") + + for relative_path in (str(item) for item in configured_refs): + suffix = Path(relative_path).suffix or ".tmp" + temp_path = storage.download_raw_to_temp( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + suffix=suffix, + temp_dir=output_dir, + ) + target_path = Path(output_dir) / relative_path + target_path.parent.mkdir(parents=True, exist_ok=True) + Path(temp_path).replace(target_path) + + +def _upload_scope_artifacts( + *, + storage: ResultStorage, + job_id: str, + raw_prefix: str | None, + workspace_dir: str, + artifact_refs: set[str], +) -> None: + if not artifact_refs: + return + add_dirs = [ + os.path.join(workspace_dir, entry) + for entry in os.listdir(workspace_dir) + if os.path.isdir(os.path.join(workspace_dir, entry)) + ] + if not add_dirs: + return + add_dir = add_dirs[0] + for relative_path in sorted(artifact_refs): + local_path = os.path.join(add_dir, relative_path) + if not os.path.isfile(local_path): + continue + storage.upload_raw_file( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + local_file_path=local_path, + ) + + +def _publish_and_mark_refined( + *, + job_result_id: str, + scope_id: str, + target: _ClaimedScopeTarget, + output: ScopeRefinementOutput, + scope_publisher: ScopePublisher, +) -> bool | None: + now = utc_now_naive() + publication_scope = DocumentPublicationScope( + user_id=target.user_id, + namespace=target.namespace, + document_id=target.document_id, + job_result_id=job_result_id, + source_file_name=target.source_file_name, + ) + with get_sync_db_context() as db: + document = db.execute( + select(Document) + .where(Document.document_id == target.document_id) + .with_for_update() + ).scalar_one_or_none() + if ( + document is None + or document.status != "active" + or document.current_job_result_id != job_result_id + ): + db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .values(state="failed", updated_at=now) + ) + return None + + scope_publisher( + db, + scope=publication_scope, + scope_id=scope_id, + chunks=output.chunks, + section_summaries=output.section_summaries, + ) + db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .values( + state=REFINED_ENRICHMENT_STATE, + refined_at=now, + updated_at=now, + ) + ) + remaining = int( + db.execute( + select(func.count()) + .select_from(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.state != REFINED_ENRICHMENT_STATE) + ).scalar_one() + ) + return remaining == 0 + + +def _finalize_all_scopes_refined( + *, + job_result_id: str, + target: _ClaimedScopeTarget, + storage: ResultStorage, + workspace_dir: str, +) -> None: + try: + _regenerate_and_upload_full_package( + job_result_id=job_result_id, + target=target, + storage=storage, + workspace_dir=workspace_dir, + ) + except Exception as exc: + logger.warning( + "[page_memory] failed to regenerate ZIP after all scopes refined " + "job_result_id={}: {}", + job_result_id, + exc, + ) + return + _mark_job_enrichment_complete( + job_id=target.job_id, + job_result_id=job_result_id, + ) + + +def _regenerate_and_upload_full_package( + *, + job_result_id: str, + target: _ClaimedScopeTarget, + storage: ResultStorage, + workspace_dir: str, +) -> None: + with get_sync_db_context() as db: + rows = ( + db.execute( + select(DocumentChunk) + .where(DocumentChunk.document_id == target.document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .order_by(DocumentChunk.sort_order, DocumentChunk.id) + ) + .scalars() + .all() + ) + chunks = [ + { + "chunk_id": row.chunk_id, + "type": row.chunk_type, + "content": row.content or "", + "path": row.source_chunk_path or "", + "metadata": dict(row.chunk_metadata or {}), + "order": row.sort_order, + } + for row in rows + ] + if not chunks: + return + + package_dir = os.path.join(workspace_dir, "_full_package", target.source_file_name) + os.makedirs(package_dir, exist_ok=True) + for relative_path in sorted(collect_referenced_artifact_refs(chunks)): + target_path = Path(package_dir) / relative_path + if target_path.is_file(): + continue + suffix = Path(relative_path).suffix or ".tmp" + temp_path = storage.download_raw_to_temp( + job_id=target.job_id, + relative_path=relative_path, + raw_prefix=target.raw_prefix, + suffix=suffix, + temp_dir=package_dir, + ) + target_path.parent.mkdir(parents=True, exist_ok=True) + Path(temp_path).replace(target_path) + + ensure_doc_nav_json( + package_dir, + chunks, + source_file_name=target.source_file_name, + overwrite=True, + ) + enrich_doc_nav_summaries( + os.path.dirname(package_dir), + source_file=target.source_file_name, + force=True, + use_llm=bool( + JobMetadataHelper.get_parsing_param( + target.job_metadata, + "summary_use_llm", + False, + ) + ), + top_summary_use_llm=bool( + JobMetadataHelper.get_parsing_param( + target.job_metadata, + "top_summary_use_llm", + True, + ) + ), + chunks=chunks, + ) + zip_path, checksum, _statistics, zip_size = ZipResultService().generate_zip_package( + job_id=target.job_id, + chunks=chunks, + add_dir=package_dir, + source_file_name=target.source_file_name, + data_id=JobMetadataHelper.get_data_id(target.job_metadata), + job_metadata=target.job_metadata, + temp_dir=workspace_dir, + ) + bundle = storage.upload( + job_id=target.job_id, + result_dir=package_dir, + zip_file_path=zip_path, + artifact_refs=set(), + ) + checksum_value = ( + str(checksum.get("value") or "") + if isinstance(checksum, dict) + else str(checksum or "") + ) + with get_sync_db_context() as db: + db.execute( + update(JobResult) + .where(JobResult.id == job_result_id) + .values( + result_s3_key=bundle.zip_key, + result_size=zip_size, + inline_payload={"checksum": checksum_value}, + ) + ) + + +def _mark_job_enrichment_complete( + *, + job_id: str, + job_result_id: str | None = None, +) -> None: + _update_job_and_result_metadata( + job_id=job_id, + job_result_id=job_result_id, + updates={"enrichment_status": "complete"}, + ) + + +def _update_job_and_result_metadata( + *, + job_id: str, + job_result_id: str | None, + updates: dict[str, Any], +) -> None: + with get_sync_db_context() as db: + job = db.execute( + select(Job).where(Job.job_id == job_id).with_for_update() + ).scalar_one_or_none() + if job is not None: + job.job_metadata = { + **dict(job.job_metadata or {}), + **updates, + } + if job_result_id: + job_result = db.execute( + select(JobResult) + .where(JobResult.id == job_result_id) + .with_for_update(of=JobResult) + ).scalar_one_or_none() + if job_result is not None: + job_result.document_metadata = { + **dict(job_result.document_metadata or {}), + **updates, + } + try: + redis_service = SyncRedisServiceFactory.get_service() + SyncJobMetadataService(redis_service).update_metadata(job_id, updates) + except Exception as exc: + logger.debug( + "[page_memory] failed to update Redis job metadata for {}: {}", + job_id, + exc, + ) + + +def _mark_scope_failed(*, job_result_id: str, scope_id: str) -> None: + now = utc_now_naive() + with get_sync_db_context() as db: + db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .values(state="failed", updated_at=now) + ) + + +def _build_scope_section_summaries( + *, + output_dir: str, + chunks: list[dict[str, Any]], + source_file_name: str, + summary_use_llm: bool, +) -> dict[str, str]: + if not chunks: + return {} + nav_dir = os.path.join(output_dir, "_scope_nav", source_file_name) + os.makedirs(nav_dir, exist_ok=True) + ensure_doc_nav_json( + nav_dir, + chunks, + source_file_name=source_file_name, + overwrite=True, + ) + enrich_doc_nav_summaries( + os.path.dirname(nav_dir), + source_file=source_file_name, + force=True, + use_llm=summary_use_llm, + top_summary_use_llm=False, + chunks=chunks, + ) + lookup = build_section_summary_lookup(nav_dir) + lookup.pop("Root", None) + return lookup + + +def _find_manifest_scope( + manifest: dict[str, Any], + scope_id: str, +) -> tuple[dict[str, Any], int, int]: + scopes = [ + entry for entry in (manifest.get("scopes") or []) if isinstance(entry, dict) + ] + for index, entry in enumerate(scopes, start=1): + if str(entry.get("scope_id") or "") == scope_id: + return entry, index, len(scopes) + raise ValueError(f"scope_id {scope_id!r} not found in phase A manifest") + + +def _load_anatomy_features_and_labels( + anatomy_path: str, +) -> tuple[list[PageFeature], list[PageLabel]]: + raw = _read_optional_json(anatomy_path) + if not isinstance(raw, dict): + return [], [] + return _parse_features_and_labels( + raw.get("page_features") or [], + raw.get("page_labels") or [], + ) + + +def _parse_features_and_labels( + raw_features: list[Any], + raw_labels: list[Any], +) -> tuple[list[PageFeature], list[PageLabel]]: + features: list[PageFeature] = [] + for item in raw_features: + if not isinstance(item, dict): + continue + orientation = item.get("orientation") + features.append( + PageFeature( + page=int(item.get("page") or 0), + raw_text_length=int(item.get("raw_text_length") or 0), + text_density=float(item.get("text_density") or 0.0), + image_coverage=float(item.get("image_coverage") or 0.0), + image_count=int(item.get("image_count") or 0), + table_count=int(item.get("table_count") or 0), + drawings_count=int(item.get("drawings_count") or 0), + orientation="landscape" if orientation == "landscape" else "portrait", + width=float(item.get("width") or 0.0), + height=float(item.get("height") or 0.0), + has_asset=bool(item.get("has_asset", False)), + is_blank_like=bool(item.get("is_blank_like", False)), + invisible_text_length=int(item.get("invisible_text_length") or 0), + ) + ) + labels: list[PageLabel] = [] + for item in raw_labels: + if not isinstance(item, dict): + continue + kind = item.get("kind") + labels.append( + PageLabel( + page=int(item.get("page") or 0), + kind="landscape" if kind == "landscape" else "normal", + evidence=dict(item.get("evidence") or {}), + ) + ) + return features, labels + + +def _read_json_dict(path: str) -> dict[str, Any]: + data = _read_optional_json(path) + if not isinstance(data, dict): + raise ValueError(f"Expected JSON object at {path}") + return data + + +def _read_optional_json(path: str) -> Any: + if not os.path.isfile(path): + return None + with open(path, encoding="utf-8") as handle: + return json.load(handle) + + +def _resolve_default_scope_publisher() -> ScopePublisher: + from shared.services.retrieval.publication_scope_content import ( + replace_document_scope_content, + ) + + return replace_document_scope_content diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 8154ad1bb..463c1001b 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -24,6 +24,7 @@ ) from app.services.document_agent.structure.toc_anchoring import ( body_pages_excluding_toc, + evidenced_physical_span, pending_toc_body_scope, select_global_toc_hierarchies, ) @@ -134,14 +135,13 @@ def extract_section_skeletons( pending_tocs=pending_tocs, pending_records=pending_records, ) - parallel_tocs = [toc for toc, _record in parallel_pending] primary_page_count = page_count primary_body_pages = body_pages - pending_starts: list[int] = [] - for toc in parallel_tocs: - start = toc_range_start(toc) - if start is not None: - pending_starts.append(start) + pending_starts = [ + start + for toc, record in parallel_pending + if (start := _parallel_forest_start_page(toc, record)) is not None + ] if pending_starts: primary_page_count = min(pending_starts) - 1 primary_body_pages = [page for page in body_pages if page <= primary_page_count] @@ -312,6 +312,20 @@ def _parallel_pending_tocs( return parallel_pending +def _parallel_forest_start_page( + pending_toc: dict[str, Any], + record: dict[str, Any], +) -> int | None: + anchor_raw = record.get("skeleton_anchor") + if isinstance(anchor_raw, dict): + span = evidenced_physical_span( + deserialize_skeleton_anchor(anchor_raw).match_overrides + ) + if span is not None: + return span[0] + return toc_range_start(pending_toc) + + def _resolve_pending_tocs( *, parallel_pending: list[tuple[dict[str, Any], dict[str, Any]]], @@ -346,12 +360,19 @@ def _resolve_pending_tocs( if offset is None: raise ValueError("pending TOC offset missing after PROFILE classify") - toc_scope_end, toc_body_pages = pending_toc_body_scope( + _, toc_body_pages = pending_toc_body_scope( pending_tocs=parallel_tocs, index=i, page_count=page_count, body_pages=body_pages, ) + next_starts = [ + start + for next_toc, next_record in parallel_pending[i + 1 :] + if (start := _parallel_forest_start_page(next_toc, next_record)) is not None + ] + toc_scope_end = (min(next_starts) - 1) if next_starts else page_count + toc_body_pages = [page for page in toc_body_pages if page <= toc_scope_end] match_overrides = skeleton_anchor.match_overrides null_page_report = skeleton_anchor.null_page_report locate_summary: dict[str, Any] = { diff --git a/apps/worker/requirements.txt b/apps/worker/requirements.txt index 8d72aa832..44df7ce58 100644 --- a/apps/worker/requirements.txt +++ b/apps/worker/requirements.txt @@ -302,7 +302,7 @@ pandas==2.3.1 # via # knowhere-worker-app # tabula-py -pgvector==0.2.4 +pgvector==0.5.1 # via knowhere-shared pillow==12.2.0 # via diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py index 144dd2837..b6d68428b 100644 --- a/apps/worker/scripts/page_memory/_debug_pm_shared.py +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -37,7 +37,6 @@ from loguru import logger from app.services.page_memory._utils import ( - derive_excluded_toc_pages, page_scope_info, scope_id_for_pages, sort_skeletons, @@ -861,35 +860,6 @@ def load_pipeline_state( return data -def _pipeline_skeleton_rows(state: dict[str, Any]) -> list[dict[str, Any]]: - stages = state.get("stages") - stage2 = stages.get("stage2") if isinstance(stages, dict) else None - rows = stage2.get("skeletons") if isinstance(stage2, dict) else None - if not isinstance(rows, list): - raise ValueError("pipeline state is missing stages.stage2.skeletons[]") - return [row for row in rows if isinstance(row, dict)] - - -def load_pipeline_skeletons( - state_path: Path, -) -> list[Any]: - from app.services.page_memory.skeleton_extractor import SectionSkeleton - - state = load_pipeline_state(state_path) - return [ - SectionSkeleton( - section_path=str(row["section_path"]), - title=str(row["title"]), - level=int(row["level"]), - start_page=int(row["start_page"]), - end_page=int(row["end_page"]), - parent_path=row.get("parent_path"), - evidence=dict(row.get("evidence") or {}), - ) - for row in _pipeline_skeleton_rows(state) - ] - - def update_pipeline_state( state_path: Path, *, @@ -1521,53 +1491,6 @@ def _serialize_skeletons(skeletons: list[Any]) -> list[dict[str, Any]]: # ── Coarse scope helpers ───────────────────────────────────────────────────── -def build_debug_coarse_scopes( - *, - skeletons: list[Any], - filename: str, - 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 - - toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None - toc_pages = derive_excluded_toc_pages( - skeletons=skeletons, - 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, - page_count=page_count, - ) - return [ - { - "scope_id": scope.scope_id, - "skeletons": scope.skeletons, - "start_page": scope.start_page, - "end_page": scope.end_page, - "strategy": scope.strategy, - "processing_pages": pages_excluding_toc( - list(range(scope.start_page, scope.end_page + 1)), - toc_pages, - ), - "excluded_toc_pages": sorted( - page - for page in toc_page_set - if scope.start_page <= page <= scope.end_page - ), - } - for scope in scopes - ] - - def add_scope_selection_args(parser: argparse.ArgumentParser) -> None: """Flags shared by stage 3/4/5 for picking one or more coarse scopes.""" parser.add_argument( diff --git a/apps/worker/scripts/page_memory/debug_pm_progressive.py b/apps/worker/scripts/page_memory/debug_pm_progressive.py new file mode 100644 index 000000000..88d6194c6 --- /dev/null +++ b/apps/worker/scripts/page_memory/debug_pm_progressive.py @@ -0,0 +1,662 @@ +#!/usr/bin/env python3 +# ruff: noqa: E402 +"""Debug the production two-level progressive ingestion path. + +Calls the same functions the worker uses. Does not reimplement profiling, +Phase A, or per-scope C1-C7. + + M0 memory_service.run(ingestion_mode="progressive") + coarse _resume_profile_from_m0_workspace + build_phase_a + persist_phase_b_inputs + fine run_single_scope_refinement + +Usage: + cd apps/worker + uv run python scripts/page_memory/debug_pm_progressive.py --file /path/doc.pdf --step m0 + uv run python scripts/page_memory/debug_pm_progressive.py --file /path/doc.pdf --step coarse + uv run python scripts/page_memory/debug_pm_progressive.py --file /path/doc.pdf --step coarse --list-scopes + uv run python scripts/page_memory/debug_pm_progressive.py --file /path/doc.pdf --step fine --list-scopes + uv run python scripts/page_memory/debug_pm_progressive.py --file /path/doc.pdf --step fine --scope-id p14-23 + uv run python scripts/page_memory/debug_pm_progressive.py --file /path/doc.pdf --step fine --all-scopes +""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path as _Path + +sys.path.insert(0, str(_Path(__file__).resolve().parent)) + +import time +from typing import Any, Literal + +from loguru import logger + +from _debug_pm_shared import ( + TokenCostTracker, + add_scope_selection_args, + base_argparser, + require_file, + resolve_paths, + stop_with_trace, + update_pipeline_state, + pipeline_state_path, + record_stage, +) + +ProgressiveStep = Literal["m0", "coarse", "fine"] +_DEBUG_STATE_NAME = "progressive_debug_state.json" + + +def main() -> int: + parser = base_argparser( + "Progressive ingestion debug: M0 text → coarse structure → per-scope fine" + ) + parser.add_argument( + "--step", + choices=("m0", "coarse", "fine"), + required=True, + help="Which production progressive step to run", + ) + add_scope_selection_args(parser) + parser.add_argument( + "--run-db", + action="store_true", + help="Publish this step to the local debug DB (same revision across steps)", + ) + parser.add_argument( + "--publish-job-id", + default=None, + help="Explicit job_id for --run-db (M0 creates it; later steps reuse it)", + ) + parser.add_argument( + "--skip-assets", + action="store_true", + help="With --run-db, skip asset upload", + ) + args = parser.parse_args() + + pdf_path, filename, out_dir = resolve_paths(args) + logger.info("█" * 70) + logger.info(f" PROGRESSIVE {args.step.upper()} — {filename}") + logger.info(f" OUTPUT: {out_dir}") + logger.info("█" * 70) + + t_start = time.time() + trace_stages: list[dict[str, Any]] = [] + token_cost_tracker = TokenCostTracker() + extra: dict[str, Any] + restore_models = _apply_model_overrides(args) + + try: + if args.step == "m0": + extra = _run_m0( + pdf_path=pdf_path, + filename=filename, + out_dir=out_dir, + args=args, + trace_stages=trace_stages, + token_cost_tracker=token_cost_tracker, + ) + elif args.step == "coarse": + extra = _run_coarse( + filename=filename, + out_dir=out_dir, + args=args, + trace_stages=trace_stages, + token_cost_tracker=token_cost_tracker, + ) + else: + extra = _run_fine( + filename=filename, + out_dir=out_dir, + args=args, + trace_stages=trace_stages, + token_cost_tracker=token_cost_tracker, + ) + finally: + restore_models() + + elapsed = time.time() - t_start + update_pipeline_state( + pipeline_state_path(out_dir), + stage={"m0": 0, "coarse": 2, "fine": 5}[args.step], + payload={"progressive_step": args.step, **extra, "elapsed_s": elapsed}, + ) + return stop_with_trace( + out_dir=out_dir, + stages=trace_stages, + stop_at=f"progressive_{args.step}", + page_count=int(extra.get("page_count") or 0), + pipeline_stage={"m0": 0, "coarse": 2, "fine": 5}[args.step], + elapsed_s=elapsed, + token_cost_tracker=token_cost_tracker, + extra_summary=extra, + ) + + +def _run_m0( + *, + pdf_path: str, + filename: str, + out_dir: _Path, + args: Any, + trace_stages: list[dict[str, Any]], + token_cost_tracker: TokenCostTracker, +) -> dict[str, Any]: + from app.services.page_memory.memory_service import PageMemoryInput, run + from app.services.page_memory.phase_a import ( + PHASE_B_BLACKBOARD_REL, + PHASE_B_BODY_TEXT_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_SOURCE_PDF_REL, + ) + from shared.models.schemas.page_memory_config import PageMemoryConfig + from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks + + page_output_dir, parsed_df = run( + PageMemoryInput( + file_path=pdf_path, + filename=filename, + output_dir=str(out_dir.parent), + job_id=filename, + internal_output_filename=out_dir.name, + page_memory_config=PageMemoryConfig.default(), + ingestion_mode="progressive", + ) + ) + if _Path(page_output_dir).resolve() != out_dir.resolve(): + raise RuntimeError( + f"M0 wrote {page_output_dir}, expected {out_dir}. " + "Do not change PageMemoryInput.output_dir / internal_output_filename." + ) + + chunks = list(dataframe_to_chunks(parsed_df)) + _write_chunk_bundle(out_dir / "m0_chunks.json", chunks) + _write_chunk_bundle(out_dir / "chunks.json", chunks) + record_stage( + trace_stages, + "progressive.m0", + page_info={"page_count": len(chunks)}, + variables={ + "chunk_count": len(chunks), + "blackboard": str(out_dir / PHASE_B_BLACKBOARD_REL), + "source_pdf": str(out_dir / PHASE_B_SOURCE_PDF_REL), + }, + ) + token_cost_tracker.snapshot_stage("progressive.m0") + _maybe_publish( + out_dir=out_dir, + filename=filename, + chunks=chunks, + enrichment_status="text", + args=args, + replace_scope_id=None, + ) + extra = { + "page_count": len(chunks), + "chunk_count": len(chunks), + "phase_b": [ + PHASE_B_SOURCE_PDF_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, + ], + } + logger.info(" M0 pages={} → {}", len(chunks), out_dir / "m0_chunks.json") + return extra + + +def _run_coarse( + *, + filename: str, + out_dir: _Path, + args: Any, + trace_stages: list[dict[str, Any]], + token_cost_tracker: TokenCostTracker, +) -> dict[str, Any]: + from app.services.page_memory.phase_a import ( + PHASE_B_ANATOMY_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_SOURCE_PDF_REL, + build_phase_a, + persist_phase_b_inputs, + ) + from app.services.page_memory.scope_enrichment_service import ( + _ClaimedScopeTarget, + _build_scope_section_summaries, + _resume_profile_from_m0_workspace, + ) + from shared.models.schemas.page_memory_config import PageMemoryConfig + from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks + + require_file( + out_dir / PHASE_B_BLACKBOARD_REL, + hint="Run --step m0 first.", + ) + require_file( + out_dir / PHASE_B_SOURCE_PDF_REL, + hint="Run --step m0 first (source.pdf is written by memory_service.run).", + ) + + page_memory_config = PageMemoryConfig.default() + profile = _resume_profile_from_m0_workspace( + add_dir=str(out_dir), + target=_ClaimedScopeTarget( + document_id="debug", + job_id=filename, + user_id="debug_local_user", + namespace="default", + source_file_name=filename, + raw_prefix=None, + job_metadata={}, + page_memory_config=page_memory_config, + ), + ) + phase_a = build_phase_a(profile=profile, filename=filename) + persist_phase_b_inputs( + output_dir=str(out_dir), + profile=profile, + manifest=phase_a.manifest, + ) + chunks = list(dataframe_to_chunks(phase_a.coarse_df)) + _write_chunk_bundle(out_dir / "coarse_chunks.json", chunks) + _write_chunk_bundle(out_dir / "chunks.json", chunks) + _build_scope_section_summaries( + output_dir=str(out_dir), + chunks=chunks, + source_file_name=filename, + summary_use_llm=False, + ) + scopes = list(phase_a.manifest.get("scopes") or []) + _print_scopes(scopes) + if args.list_scopes: + logger.info(" --list-scopes: coarse built, printed scopes, not publishing") + record_stage( + trace_stages, + "progressive.coarse", + page_info={"page_count": int(phase_a.manifest.get("page_count") or 0)}, + variables={ + "scope_count": len(scopes), + "chunk_count": len(chunks), + "manifest": str(out_dir / PHASE_B_MANIFEST_REL), + "anatomy": str(out_dir / PHASE_B_ANATOMY_REL), + }, + ) + token_cost_tracker.snapshot_stage("progressive.coarse") + if not args.list_scopes: + _maybe_publish( + out_dir=out_dir, + filename=filename, + chunks=chunks, + enrichment_status="coarse", + args=args, + replace_scope_id="*", + write_scope_rows=True, + ) + extra = { + "page_count": int(phase_a.manifest.get("page_count") or 0), + "chunk_count": len(chunks), + "scope_ids": [str(entry.get("scope_id") or "") for entry in scopes], + } + logger.info( + " coarse scopes={} chunks={} → {}", + len(scopes), + len(chunks), + out_dir / "coarse_chunks.json", + ) + return extra + + +def _run_fine( + *, + filename: str, + out_dir: _Path, + args: Any, + trace_stages: list[dict[str, Any]], + token_cost_tracker: TokenCostTracker, +) -> dict[str, Any]: + from app.services.document_agent.pdf_text import coerce_page_text_cache + from app.services.page_memory.phase_a import ( + PHASE_B_ANATOMY_REL, + PHASE_B_BODY_TEXT_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_SOURCE_PDF_REL, + ) + from app.services.page_memory.scope_enrichment_service import ( + _load_anatomy_features_and_labels, + _read_json_dict, + run_single_scope_refinement, + ) + from shared.models.schemas.page_memory_config import PageMemoryConfig + + require_file(out_dir / PHASE_B_MANIFEST_REL, hint="Run --step coarse first.") + require_file(out_dir / PHASE_B_SOURCE_PDF_REL, hint="Run --step m0 first.") + require_file(out_dir / PHASE_B_FULL_TEXT_REL, hint="Run --step m0 first.") + require_file(out_dir / PHASE_B_ANATOMY_REL, hint="Run --step coarse first.") + + manifest = _read_json_dict(str(out_dir / PHASE_B_MANIFEST_REL)) + scopes = [entry for entry in (manifest.get("scopes") or []) if isinstance(entry, dict)] + _print_scopes(scopes) + if args.list_scopes: + return { + "page_count": int(manifest.get("page_count") or 0), + "scope_ids": [str(entry.get("scope_id") or "") for entry in scopes], + } + + selected = _selected_scope_ids(args=args, scopes=scopes) + page_texts = coerce_page_text_cache( + _read_json_dict(str(out_dir / PHASE_B_FULL_TEXT_REL)) + ) + body_text_by_page = coerce_page_text_cache( + _read_json_dict(str(out_dir / PHASE_B_BODY_TEXT_REL)) + if (out_dir / PHASE_B_BODY_TEXT_REL).is_file() + else {} + ) + page_features, page_labels = _load_anatomy_features_and_labels( + str(out_dir / PHASE_B_ANATOMY_REL) + ) + page_memory_config = PageMemoryConfig.default() + refined_scope_ids: list[str] = [] + last_chunk_count = 0 + for scope_id in selected: + output = run_single_scope_refinement( + scope_id=scope_id, + pdf_path=str(out_dir / PHASE_B_SOURCE_PDF_REL), + output_dir=str(out_dir), + manifest=manifest, + page_texts=page_texts, + body_text_by_page=body_text_by_page, + page_features=page_features, + page_labels=page_labels, + page_memory_config=page_memory_config, + source_file_name=filename, + summary_use_llm=False, + ) + scope_dir = out_dir / "scopes" / scope_id + scope_dir.mkdir(parents=True, exist_ok=True) + _write_chunk_bundle(scope_dir / "refined_chunks.json", output.chunks) + last_chunk_count = len(output.chunks) + refined_scope_ids.append(scope_id) + record_stage( + trace_stages, + "progressive.fine", + page_info={"scope_id": scope_id}, + variables={ + "scope_id": scope_id, + "chunk_count": len(output.chunks), + "artifact_refs": sorted(output.artifact_refs), + }, + ) + token_cost_tracker.snapshot_stage(f"progressive.fine.{scope_id}") + _maybe_publish( + out_dir=out_dir, + filename=filename, + chunks=output.chunks, + enrichment_status="coarse", + args=args, + replace_scope_id=scope_id, + ) + logger.info( + " fine {} chunks={} → {}", + scope_id, + len(output.chunks), + scope_dir / "refined_chunks.json", + ) + return { + "page_count": int(manifest.get("page_count") or 0), + "chunk_count": last_chunk_count, + "scope_ids": refined_scope_ids, + } + + +def _apply_model_overrides(args: Any): + from shared.core.config import settings + + previous_image = settings.IMAGE_MODEL + previous_hierarchy = settings.HIERARCHY_LLM_MODEL + previous_env = os.environ.get("IMAGE_MODEL") + if args.vlm_model: + settings.IMAGE_MODEL = args.vlm_model + os.environ["IMAGE_MODEL"] = args.vlm_model + if args.model: + settings.HIERARCHY_LLM_MODEL = args.model + + def restore() -> None: + settings.IMAGE_MODEL = previous_image + settings.HIERARCHY_LLM_MODEL = previous_hierarchy + if previous_env is None: + os.environ.pop("IMAGE_MODEL", None) + else: + os.environ["IMAGE_MODEL"] = previous_env + + return restore + + +def _selected_scope_ids(*, args: Any, scopes: list[dict[str, Any]]) -> list[str]: + available_rows = [ + { + "scope_id": str(entry.get("scope_id") or ""), + "start_page": int(entry.get("start_page") or 0), + "end_page": int(entry.get("end_page") or 0), + } + for entry in scopes + if entry.get("scope_id") + ] + available = [row["scope_id"] for row in available_rows] + raw = str(getattr(args, "scope_id", "") or "").strip() + if raw: + wanted = [part.strip() for part in raw.split(",") if part.strip()] + missing = [scope_id for scope_id in wanted if scope_id not in available] + if missing: + raise SystemExit( + f"Unknown --scope-id {missing}. Available: {', '.join(available)}" + ) + return wanted + page_range = str(getattr(args, "page_range", "") or "").strip() + if page_range: + parts = page_range.split("-") + try: + start = int(parts[0]) + end = int(parts[1]) if len(parts) > 1 else start + except ValueError: + raise SystemExit(f"Invalid --page-range {page_range!r}; expected e.g. 14-23") + selected = [ + row["scope_id"] + for row in available_rows + if row["start_page"] <= end and row["end_page"] >= start + ] + if not selected: + raise SystemExit( + f"No scope overlaps --page-range {start}-{end}. Available: {', '.join(available)}" + ) + return selected + if args.all_scopes or args.fat_only: + if args.fat_only and available_rows: + fattest = max( + available_rows, + key=lambda row: int(row["end_page"]) - int(row["start_page"]), + ) + return [fattest["scope_id"]] + return available + raise SystemExit( + "fine step needs --scope-id, --page-range, --all-scopes, --fat-only, or --list-scopes" + ) + + +def _print_scopes(scopes: list[dict[str, Any]]) -> None: + logger.info(" scopes ({}):", len(scopes)) + for entry in scopes: + logger.info( + " {} p{}-{} strategy={}", + entry.get("scope_id"), + entry.get("start_page"), + entry.get("end_page"), + entry.get("strategy"), + ) + + +def _write_chunk_bundle(path: _Path, chunks: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps({"chunks": chunks}, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + +def _maybe_publish( + *, + out_dir: _Path, + filename: str, + chunks: list[dict[str, Any]], + enrichment_status: str, + args: Any, + replace_scope_id: str | None, + write_scope_rows: bool = False, +) -> None: + if not args.run_db: + return + from scripts._debug_publish import publish_debug_result_dir + from shared.core.database_sync import get_sync_db_context + from shared.models.database.document import Document + from shared.models.database.job import Job + from shared.models.database.job_result import JobResult + from sqlalchemy import select + + state = _load_debug_state(out_dir) + if state is None or replace_scope_id is None: + from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder + + doc_nav = ZipDocNavigationBuilder().build_doc_nav(chunks, filename) + (out_dir / "doc_nav.json").write_text( + json.dumps(doc_nav, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + published = publish_debug_result_dir( + result_dir=out_dir, + source_file_name=filename, + chunks=chunks, + job_id=args.publish_job_id, + parse_track="page_memory", + upload_assets=not args.skip_assets, + ) + with get_sync_db_context() as db: + document = db.execute( + select(Document).where(Document.document_id == published.document_id) + ).scalar_one() + job = db.execute( + select(Job).where(Job.job_id == published.job_id) + ).scalar_one() + job.job_metadata = { + **dict(job.job_metadata or {}), + "parse_track": "page_memory", + "ingestion_mode": "progressive", + "enrichment_status": enrichment_status, + } + if document.current_job_result_id: + job_result = db.execute( + select(JobResult).where( + JobResult.id == document.current_job_result_id + ) + ).scalar_one() + job_result.document_metadata = { + **dict(job_result.document_metadata or {}), + "parse_track": "page_memory", + "ingestion_mode": "progressive", + "enrichment_status": enrichment_status, + } + state = { + "job_id": published.job_id, + "document_id": published.document_id, + "job_result_id": document.current_job_result_id, + "user_id": published.user_id, + "namespace": published.namespace, + } + if write_scope_rows and published.document_id and document.current_job_result_id: + from shared.services.jobs.lifecycle.scope_enrichments import ( + write_coarse_scope_enrichments, + ) + + write_coarse_scope_enrichments( + db, + document_id=published.document_id, + job_result_id=document.current_job_result_id, + ) + _write_debug_state(out_dir, state) + logger.info( + " published job_id={} document_id={} status={}", + state["job_id"], + state["document_id"], + enrichment_status, + ) + return + + from shared.services.jobs.lifecycle.scope_enrichments import ( + write_coarse_scope_enrichments, + ) + from shared.services.retrieval.publication_models import DocumentPublicationScope + from shared.services.retrieval.publication_scope_content import ( + replace_document_scope_content, + ) + + with get_sync_db_context() as db: + replace_document_scope_content( + db, + scope=DocumentPublicationScope( + user_id=str(state["user_id"]), + namespace=str(state["namespace"]), + document_id=str(state["document_id"]), + job_result_id=str(state["job_result_id"]), + source_file_name=filename, + ), + scope_id=replace_scope_id, + chunks=chunks, + ) + if write_scope_rows: + write_coarse_scope_enrichments( + db, + document_id=str(state["document_id"]), + job_result_id=str(state["job_result_id"]), + ) + job = db.execute(select(Job).where(Job.job_id == state["job_id"])).scalar_one() + job.job_metadata = { + **dict(job.job_metadata or {}), + "enrichment_status": enrichment_status, + } + job_result = db.execute( + select(JobResult).where(JobResult.id == state["job_result_id"]) + ).scalar_one() + job_result.document_metadata = { + **dict(job_result.document_metadata or {}), + "enrichment_status": enrichment_status, + } + logger.info( + " replaced scope={} on document_id={} status={}", + replace_scope_id, + state["document_id"], + enrichment_status, + ) + + +def _load_debug_state(out_dir: _Path) -> dict[str, Any] | None: + path = out_dir / _DEBUG_STATE_NAME + if not path.is_file(): + return None + raw = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(raw, dict) or not raw.get("job_result_id"): + return None + return raw + + +def _write_debug_state(out_dir: _Path, state: dict[str, Any]) -> None: + (out_dir / _DEBUG_STATE_NAME).write_text( + json.dumps(state, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py b/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py index 36a448a63..7035eebf8 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py @@ -7,8 +7,9 @@ restore Stage-1 outline (or re-probe) → calibrate / outline route → classify contained/parallel → graft contained → write skeleton_* -Also resolves coarse skeletons (C4 resolve-only) into pipeline state so -Stage 3 can resume without re-anchoring. No scope build or fine hierarchy. +Also resolves coarse skeletons (C4 resolve-only) and the Phase A scope +manifest into pipeline state so Stage 3 can resume without re-anchoring +or rebuilding scopes. No fine hierarchy. Requires Stage 0 → Stage 1 first: uv run python scripts/page_memory/debug_pm_stage0_bootstrap.py --file ... @@ -75,7 +76,7 @@ def main() -> int: from app.services.document_agent.persist import DOC_PROFILE_FILENAME from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring from app.services.document_agent.validators import single_shard_plan - from app.services.page_memory.skeleton_extractor import extract_section_skeletons + from app.services.page_memory.phase_a import build_phase_a from shared.core.config import settings pdf_path, filename, out_dir = resolve_paths(args) @@ -143,13 +144,20 @@ def main() -> int: pass from app.services.document_agent.pdf_text import page_content_map + from types import SimpleNamespace page_texts = page_content_map(bb.page_full_text_cache or {}) - skeletons = extract_section_skeletons( - anatomy=anchored, + phase_a = build_phase_a( + profile=SimpleNamespace( + page_count=page_count, + anatomy=anchored, + page_full_text_cache=page_texts, + page_body_text_cache={}, + ), filename=filename, - page_texts=page_texts, ) + skeletons = phase_a.skeletons + scope_manifest = phase_a.manifest finally: if args.model: settings.IMAGE_MODEL = previous_image_model @@ -177,6 +185,7 @@ def main() -> int: "pending_skeleton_anchors": pending_records, "pending_summary": pending_summary, "skeletons": _serialize_skeletons(skeletons), + "scope_manifest": scope_manifest, }, ) @@ -222,7 +231,11 @@ def main() -> int: pipeline_stage=2, elapsed_s=elapsed, token_cost_tracker=token_cost_tracker, - extra_summary={"pending": pending_summary, "skeleton_count": len(skeletons)}, + extra_summary={ + "pending": pending_summary, + "skeleton_count": len(skeletons), + "scope_count": len(scope_manifest.get("scopes") or []), + }, ) 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 ce482cd60..fa92b3fff 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 @@ -5,8 +5,9 @@ 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``). +Requires Stage 2 output: _doc_agent/pipeline_state.json (with scope_manifest), +doc_profile.json (after production ``run_toc_anchoring``). Stage 3 reads the +frozen Phase A manifest and does not rebuild scopes. Usage: cd apps/worker @@ -34,10 +35,9 @@ TraceStageAdapter, add_scope_selection_args, base_argparser, - build_debug_coarse_scopes, load_anatomy_cache, resolve_anatomy_cache_path, - load_pipeline_skeletons, + load_pipeline_state, load_scope_skeletons_artifact, list_scope_dirs, pipeline_state_path, @@ -67,9 +67,8 @@ def _resolve_scope_processing_pages( toc_pages: list[int], ) -> tuple[list[int], list[int]]: from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc - from app.services.page_memory._utils import derive_excluded_toc_pages - toc_pages = derive_excluded_toc_pages(skeletons=skeletons, toc_pages=toc_pages) + toc_pages = sorted({int(page) for page in (toc_pages or [])}) processing_pages = [ int(page) for page in (scope_meta.get("processing_pages") or []) ] @@ -134,8 +133,7 @@ def _run_fine_hierarchy_for_scope( 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) - from app.services.page_memory._utils import derive_excluded_toc_pages - toc_pages = derive_excluded_toc_pages(skeletons=active_skeletons, toc_pages=toc_pages) + toc_pages = sorted({int(page) for page in (toc_pages or [])}) strategy = str(scope_meta.get("strategy") or "coarse_scope") processing_pages, excluded_toc_pages = _resolve_scope_processing_pages( scope_meta=scope_meta, @@ -336,8 +334,12 @@ def main() -> int: 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.skeleton_extractor import SectionSkeleton + from app.services.page_memory.phase_a import ( + next_title_by_path_from_manifest, + require_phase_a_manifest, + scopes_from_phase_a_manifest, + skeletons_from_phase_a_manifest, + ) from shared.models.schemas.page_memory_config import PageMemoryConfig pdf_path, filename, out_dir = resolve_paths(args) @@ -358,12 +360,19 @@ 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_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) - from app.services.page_memory._utils import derive_excluded_toc_pages - toc_pages = derive_excluded_toc_pages(skeletons=skeletons, toc_pages=toc_pages) + stage2 = load_pipeline_state(state_path).get("stages", {}).get("stage2") + try: + manifest = require_phase_a_manifest(stage2 if isinstance(stage2, dict) else None) + except ValueError as exc: + raise SystemExit( + f"{exc}. Re-run Stage 2: " + "uv run python scripts/page_memory/debug_pm_stage2_calibration.py --file ..." + ) from exc + skeletons = skeletons_from_phase_a_manifest(manifest) + toc_pages = [int(page) for page in (manifest.get("toc_pages") or [])] + coarse_scopes = scopes_from_phase_a_manifest(manifest) + next_title_by_path = next_title_by_path_from_manifest(manifest) logger.info("█" * 70) logger.info(f" STAGE 3: SCOPE + FINE HIERARCHY — {filename}") @@ -374,41 +383,11 @@ def main() -> int: trace_stages: list[dict] = [] token_cost_tracker = TokenCostTracker() - # ── Build coarse scopes ── - coarse_scopes = build_debug_coarse_scopes( - skeletons=skeletons, - filename=filename, - page_count=page_count, - anatomy=anatomy, - ) - if not coarse_scopes: - root_skel = SectionSkeleton( - section_path=f"{filename}/Root", - level=1, - start_page=1, - end_page=page_count, - title="Root", - parent_path=filename, - evidence={"source": "fallback_root"}, + raise SystemExit( + "Stage 2 scope_manifest has no scopes. Re-run Stage 2: " + "uv run python scripts/page_memory/debug_pm_stage2_calibration.py --file ..." ) - coarse_scopes = [ - { - "scope_id": scope_id_for_pages(1, page_count), - "skeletons": [root_skel], - "start_page": 1, - "end_page": page_count, - "strategy": "fallback_root", - "processing_pages": pages_excluding_toc( - list(range(1, page_count + 1)), - toc_pages, - ), - "excluded_toc_pages": sorted( - {int(page) for page in toc_pages} - ), - } - ] - logger.info(" no skeleton hierarchy → fallback Root scope p1-{}", page_count) if args.list_scopes: logger.info("Available scopes ({}):", len(coarse_scopes)) @@ -541,7 +520,6 @@ def main() -> int: if partial_run: logger.info(" MODE: partial — will not overwrite top-level hierarchy.json") - next_title_by_path = build_next_title_by_path(skeletons) logger.info( " next_title_by_path: {} paths ({} with tail anchor)", len(next_title_by_path), 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 c08f0c0d9..f0a0586f7 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py @@ -74,8 +74,7 @@ def _load_scope_asset_context( hint=f"Run Stage 3 first to produce {fine_hierarchy_path}", ) prior_scope, active_skeletons = load_hierarchy_artifact(fine_hierarchy_path) - from app.services.page_memory._utils import derive_excluded_toc_pages - toc_pages = derive_excluded_toc_pages(skeletons=active_skeletons, toc_pages=toc_pages) + toc_pages = sorted({int(page) for page in (toc_pages or [])}) if not active_skeletons: logger.warning( " [scope {}] no skeletons in fine_hierarchy.json — skipping", @@ -124,7 +123,6 @@ def main() -> int: from app.services.document_agent.pdf_text import read_page_texts from app.services.page_memory.memory_service import ( _project_assets_for_pages, - _resolve_asset_max_pages, _select_rendered_pages_with_assets, ) from app.services.page_memory.page_assets import ( @@ -215,15 +213,13 @@ def main() -> int: or os.environ.get("IMAGE_MODEL") ) pm_config = PageMemoryConfig.default() - asset_max_pages = _resolve_asset_max_pages(page_count, pm_config) asset_rendered = _select_rendered_pages_with_assets(rendered, page_features) summary_enabled = page_asset_summary_enabled() logger.info( " C5: {}/{} rendered pages have coarse has_asset " - "(asset_max_pages={} summary_enabled={} model={})", + "(summary_enabled={} model={})", len(asset_rendered), len(rendered), - asset_max_pages, summary_enabled, asset_model, ) @@ -233,7 +229,6 @@ def main() -> int: rendered_pages=asset_rendered, output_dir=str(out_dir), model_name=asset_model, - max_pages=asset_max_pages, confidence_threshold=get_asset_confidence_threshold(), summary_enabled=summary_enabled, summary_concurrency=pm_config.asset_summary_concurrency, 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 2dc5c5063..bf34f8d39 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 @@ -83,8 +83,7 @@ def _run_tagging_for_scope( fine_hierarchy_path = scope_dir / "fine_hierarchy.json" require_file(fine_hierarchy_path, hint=f"Run Stage 3 to produce {fine_hierarchy_path}") prior_scope, active_skeletons = load_hierarchy_artifact(fine_hierarchy_path) - from app.services.page_memory._utils import derive_excluded_toc_pages - toc_pages = derive_excluded_toc_pages(skeletons=active_skeletons, toc_pages=toc_pages) + toc_pages = sorted({int(page) for page in (toc_pages or [])}) if not active_skeletons: logger.warning(" [scope {}] no skeletons — skipping", scope_id) return ScopeResult( diff --git a/apps/worker/tests/contract/README.md b/apps/worker/tests/contract/README.md index d9f009a5c..0ee18b4ca 100644 --- a/apps/worker/tests/contract/README.md +++ b/apps/worker/tests/contract/README.md @@ -27,7 +27,7 @@ These tests should avoid: ## Command -- `uv run python apps/api/scripts/ensure_test_environment.py --install` +- `uv run python apps/api/scripts/ensure_test_environment.py --install` (includes pgvector) - `uv run python apps/api/scripts/ensure_test_environment.py` - `uv run pytest apps/worker/tests/contract -q` diff --git a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py index 9fb278d69..f6880a7c7 100644 --- a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py +++ b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py @@ -17,11 +17,10 @@ from app.services.page_memory.page_renderer import PageRenderResult -def test_page_assets_default_to_image_model_and_full_page_scan() -> None: +def test_page_assets_default_to_image_model() -> None: from shared.models.schemas.page_memory_config import resolve_asset_model assert page_assets.get_asset_model() == resolve_asset_model() - assert page_assets.get_asset_max_pages(301) == 301 def test_page_asset_detection_uses_lowest_temperature(monkeypatch, tmp_path) -> None: @@ -144,7 +143,6 @@ def _fake_extract_table(*, asset, pdf_path, output_dir, table_engine="tabula"): rendered_pages=page_images, output_dir=str(tmp_path), model_name=page_assets.get_asset_model(), - max_pages=2, confidence_threshold=0.3, summary_enabled=False, summary_concurrency=4, diff --git a/apps/worker/tests/contract/test_page_memory_output_contract.py b/apps/worker/tests/contract/test_page_memory_output_contract.py index 3bd7684f8..aea122c74 100644 --- a/apps/worker/tests/contract/test_page_memory_output_contract.py +++ b/apps/worker/tests/contract/test_page_memory_output_contract.py @@ -238,13 +238,13 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: def test_page_memory_granularity_routes_supported_page_modes() -> None: assert ( memory_service._decide_granularity( # noqa: SLF001 - SimpleNamespace(page_count=6, toc=SimpleNamespace(has_toc=False)) + SimpleNamespace(page_count=5, toc=SimpleNamespace(has_toc=False)) ) == "whole_doc" ) assert ( memory_service._decide_granularity( # noqa: SLF001 - SimpleNamespace(page_count=7, toc=SimpleNamespace(has_toc=False)) + SimpleNamespace(page_count=6, toc=SimpleNamespace(has_toc=False)) ) == "page" ) diff --git a/apps/worker/tests/contract/test_page_memory_phase_a_contract.py b/apps/worker/tests/contract/test_page_memory_phase_a_contract.py new file mode 100644 index 000000000..2115bf4ac --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_phase_a_contract.py @@ -0,0 +1,614 @@ +"""Phase A I/O: coarse rows, frozen manifest, and result-upload handoff.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_ingestion import ( + success_finalization as success_finalization_mod, +) # noqa: E402 +from app.services.document_ingestion.parse_result_package import ( # noqa: E402 + GeneratedResultPackage, + ParseArtifact, + ParseResultPackage, +) +from app.services.page_memory.phase_a import ( # noqa: E402 + PHASE_B_ANATOMY_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_BODY_TEXT_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_SOURCE_PDF_REL, + build_coarse_text_dataframe, + build_phase_a, + next_title_by_path_from_manifest, + persist_phase_b_inputs, + require_phase_a_manifest, + scopes_from_phase_a_manifest, +) +from app.services.page_memory.skeleton_extractor import SectionSkeleton # noqa: E402 +from shared.models.schemas.page_memory_config import PageMemoryConfig # noqa: E402 +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks # noqa: E402 +from shared.services.storage.result_storage import JobResultStorage # noqa: E402 + + +def _anatomy(*, toc_pages: list[int] | None = None) -> SimpleNamespace: + return SimpleNamespace( + page_features=[], + page_labels=[], + toc_result=SimpleNamespace(toc_pages=toc_pages or []), + to_dict=lambda: {"page_count": 4, "toc_pages": toc_pages or []}, + ) + + +def _profile( + *, + skeletons: list[SectionSkeleton], + page_texts: dict[int, str], + body_texts: dict[int, str], + toc_pages: list[int] | None = None, + monkeypatch: pytest.MonkeyPatch, +) -> SimpleNamespace: + from app.services.page_memory import skeleton_extractor + + monkeypatch.setattr( + skeleton_extractor, + "extract_section_skeletons", + lambda **kwargs: skeletons, + ) + return SimpleNamespace( + page_count=4, + anatomy=_anatomy(toc_pages=toc_pages), + page_full_text_cache=page_texts, + page_body_text_cache=body_texts, + ) + + +def test_phase_a_without_anatomy_falls_back_to_root_spanning_all_pages() -> None: + profile = SimpleNamespace( + page_count=5, + anatomy=None, + page_full_text_cache={page: f"text {page}" for page in range(1, 6)}, + page_body_text_cache={}, + ) + + phase_a = build_phase_a( + profile=profile, + filename="atlas.pdf", + ) + + assert [(s.title, s.start_page, s.end_page) for s in phase_a.skeletons] == [ + ("Root", 1, 5) + ] + assert [scope.scope_id for scope in phase_a.coarse_scopes] == ["p1-5"] + + +def test_phase_a_writes_coarse_rows_and_frozen_manifest( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + skeletons = [ + SectionSkeleton("doc.pdf/A", 1, 1, 2, "A", "doc.pdf"), + SectionSkeleton("doc.pdf/B", 1, 3, 4, "B", "doc.pdf"), + ] + page_texts = {1: "page one", 2: "page two", 3: "page three", 4: "page four"} + body_texts = {1: "one", 2: "two", 3: "three", 4: "four"} + profile = _profile( + skeletons=skeletons, + page_texts=page_texts, + body_texts=body_texts, + monkeypatch=monkeypatch, + ) + + phase_a = build_phase_a( + profile=profile, + filename="doc.pdf", + ) + (tmp_path / PHASE_B_SOURCE_PDF_REL).write_bytes(b"%PDF") + written = persist_phase_b_inputs( + output_dir=str(tmp_path), + profile=profile, + manifest=phase_a.manifest, + ) + + assert [scope.scope_id for scope in phase_a.coarse_scopes] == ["p1-2", "p3-4"] + assert list(phase_a.coarse_df["path"]) == ["doc.pdf/A", "doc.pdf/B"] + assert list(phase_a.coarse_df["content"]) == [ + "page one\n\npage two", + "page three\n\npage four", + ] + extra = list(phase_a.coarse_df["extra_metadata"]) + assert extra[0]["scope_id"] == "p1-2" + assert extra[1]["scope_id"] == "p3-4" + assert extra[0]["enrichment_state"] == "coarse" + assert extra[0]["body_text_preview"] == "one two" + assert extra[1]["body_text_preview"] == "three four" + + chunks = dataframe_to_chunks(phase_a.coarse_df) + assert [chunk["metadata"]["scope_id"] for chunk in chunks] == ["p1-2", "p3-4"] + assert {chunk["metadata"]["enrichment_state"] for chunk in chunks} == {"coarse"} + assert chunks[0]["metadata"]["page_nums"] == [1, 2] + assert chunks[0]["content"] == "page one\n\npage two" + + manifest_path = tmp_path / PHASE_B_MANIFEST_REL + anatomy_path = tmp_path / PHASE_B_ANATOMY_REL + full_text_path = tmp_path / PHASE_B_FULL_TEXT_REL + body_text_path = tmp_path / PHASE_B_BODY_TEXT_REL + assert manifest_path.is_file() + assert anatomy_path.is_file() + assert full_text_path.is_file() + assert body_text_path.is_file() + on_disk = json.loads(manifest_path.read_text(encoding="utf-8")) + assert on_disk == phase_a.manifest + assert [entry["scope_id"] for entry in on_disk["scopes"]] == ["p1-2", "p3-4"] + assert on_disk["scopes"][0]["next_title_slice"] == {"doc.pdf/A": "B"} + assert on_disk["scopes"][0]["processing_pages"] == [1, 2] + assert json.loads(full_text_path.read_text(encoding="utf-8"))["1"] == "page one" + assert json.loads(body_text_path.read_text(encoding="utf-8"))["3"] == "three" + assert set(written) == { + PHASE_B_SOURCE_PDF_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_ANATOMY_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, + } + + restored = scopes_from_phase_a_manifest( + require_phase_a_manifest( + { + "scope_manifest": on_disk, + } + ) + ) + assert [scope["scope_id"] for scope in restored] == ["p1-2", "p3-4"] + assert restored[0]["skeletons"][0].title == "A" + assert next_title_by_path_from_manifest(on_disk)["doc.pdf/A"] == "B" + + +def test_coarse_rows_omit_excluded_toc_pages( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + del tmp_path + skeletons = [SectionSkeleton("doc.pdf/A", 1, 1, 4, "A", "doc.pdf")] + page_texts = {1: "page one", 2: "toc text", 3: "page three", 4: "page four"} + profile = _profile( + skeletons=skeletons, + page_texts=page_texts, + body_texts=page_texts, + toc_pages=[2], + monkeypatch=monkeypatch, + ) + phase_a = build_phase_a( + profile=profile, + filename="doc.pdf", + ) + chunks = dataframe_to_chunks(phase_a.coarse_df) + assert len(chunks) == 1 + assert chunks[0]["metadata"]["page_nums"] == [1, 3, 4] + assert chunks[0]["content"] == "page one\n\npage three\n\npage four" + assert phase_a.manifest["scopes"][0]["processing_pages"] == [1, 3, 4] + assert phase_a.manifest["scopes"][0]["excluded_toc_pages"] == [2] + + +def test_progressive_run_returns_m0_text_baseline_without_toc_or_scope_vlm( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from app.services.page_memory import memory_service, page_tagger + from app.services.page_memory.memory_service import PageMemoryInput + + skeletons = [ + SectionSkeleton("doc.pdf/A", 1, 1, 2, "A", "doc.pdf"), + SectionSkeleton("doc.pdf/B", 1, 3, 4, "B", "doc.pdf"), + ] + page_texts = {1: "one", 2: "two", 3: "three", 4: "four"} + profile = _profile( + skeletons=skeletons, + page_texts=page_texts, + body_texts=page_texts, + monkeypatch=monkeypatch, + ) + profile.page_count = 4 + profile.blackboard = SimpleNamespace( + page_count=4, + document_profile=None, + page_features=[], + page_labels=[], + doc_stats={}, + global_signals={"assets_probed": True}, + page_full_text_cache=page_texts, + ) + + source = tmp_path / "doc.pdf" + source.write_bytes(b"%PDF") + captured_stop_after: list[object] = [] + + def fail_if_tagged(**kwargs: object) -> list[object]: + raise AssertionError(f"M0 must not tag pages: {kwargs.keys()}") + + def fake_profile(*args: object, **kwargs: object) -> SimpleNamespace: + del args + captured_stop_after.append(kwargs.get("stop_after")) + return profile + + monkeypatch.setattr( + memory_service, + "normalize_to_pdf", + lambda **kwargs: (str(source), "doc.pdf"), + ) + monkeypatch.setattr(memory_service, "profile_document", fake_profile) + monkeypatch.setattr(page_tagger, "tag_pages", fail_if_tagged) + monkeypatch.setattr(page_tagger, "tag_page_titles", fail_if_tagged) + + output_dir, dataframe = memory_service.run( + PageMemoryInput( + file_path=str(source), + filename="doc.pdf", + output_dir=str(tmp_path / "out"), + page_memory_config=PageMemoryConfig(asset_extraction_enabled=False), + ingestion_mode="progressive", + ) + ) + assert captured_stop_after == ["asset_probe"] + chunks = dataframe_to_chunks(dataframe) + assert [chunk["metadata"]["enrichment_state"] for chunk in chunks] == [ + "text", + "text", + "text", + "text", + ] + assert [chunk["metadata"]["scope_id"] for chunk in chunks] == [ + "p1-1", + "p2-2", + "p3-3", + "p4-4", + ] + assert [chunk["path"] for chunk in chunks] == [ + "doc.pdf/p1", + "doc.pdf/p2", + "doc.pdf/p3", + "doc.pdf/p4", + ] + assert (Path(output_dir) / PHASE_B_BLACKBOARD_REL).is_file() + assert (Path(output_dir) / PHASE_B_FULL_TEXT_REL).is_file() + assert (Path(output_dir) / PHASE_B_SOURCE_PDF_REL).is_file() + assert not (Path(output_dir) / PHASE_B_MANIFEST_REL).exists() + + +def test_full_mode_run_invokes_full_page_dataframe_and_writes_no_phase_b_manifest( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from app.services.page_memory import memory_service + from app.services.page_memory.memory_service import PageMemoryInput + + skeletons = [SectionSkeleton("doc.pdf/A", 1, 1, 8, "A", "doc.pdf")] + page_texts = {i: f"p{i}" for i in range(1, 9)} + profile = _profile( + skeletons=skeletons, + page_texts=page_texts, + body_texts=page_texts, + monkeypatch=monkeypatch, + ) + profile.page_count = 8 + profile.toc = SimpleNamespace(has_toc=True) + + source = tmp_path / "doc.pdf" + source.write_bytes(b"%PDF") + called_full: list[bool] = [] + captured_stop_after: list[object] = [] + + def fake_profile(*args: object, **kwargs: object) -> SimpleNamespace: + del args + captured_stop_after.append(kwargs.get("stop_after")) + return profile + + def fake_build_page_df(**kwargs: object) -> pd.DataFrame: + del kwargs + called_full.append(True) + return build_coarse_text_dataframe( + page_texts={1: "full row"}, + body_text_by_page={1: "full row"}, + page_count=1, + filename="doc.pdf", + ) + + monkeypatch.setattr( + memory_service, + "normalize_to_pdf", + lambda **kwargs: (str(source), "doc.pdf"), + ) + monkeypatch.setattr(memory_service, "profile_document", fake_profile) + monkeypatch.setattr( + memory_service, "_build_page_dataframe", fake_build_page_df + ) + + output_dir, _df = memory_service.run( + PageMemoryInput( + file_path=str(source), + filename="doc.pdf", + output_dir=str(tmp_path / "out_full"), + page_memory_config=PageMemoryConfig(asset_extraction_enabled=False), + ingestion_mode="full", + ) + ) + assert captured_stop_after == [None] + assert called_full == [True] + assert not (Path(output_dir) / PHASE_B_MANIFEST_REL).exists() + assert not (Path(output_dir) / PHASE_B_BLACKBOARD_REL).exists() + + +def test_page_memory_upload_m0_marks_text_status( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + result_dir = tmp_path / "result_m0" + result_dir.mkdir() + (result_dir / "phase_b").mkdir() + (result_dir / PHASE_B_SOURCE_PDF_REL).write_bytes(b"%PDF") + (result_dir / PHASE_B_BLACKBOARD_REL).write_text("{}", encoding="utf-8") + (result_dir / PHASE_B_FULL_TEXT_REL).write_text("{}", encoding="utf-8") + (result_dir / PHASE_B_BODY_TEXT_REL).write_text("{}", encoding="utf-8") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + + class FakeStorage: + def upload(self, **kwargs: object) -> SimpleNamespace: + refs = sorted(kwargs["artifact_refs"] or []) + return SimpleNamespace( + zip_key="results/job-m0/result.zip", + raw_prefix="results/job-m0/", + raw_files={ref: f"results/job-m0/{ref}" for ref in refs}, + ) + + persisted: dict[str, object] = {} + monkeypatch.setattr( + success_finalization_mod, + "persist_job_metadata_updates", + lambda **kw: persisted.update(kw["metadata_updates"]), # type: ignore[arg-type] + ) + success_finalization_mod._upload_result_package( # noqa: SLF001 + result_package=ParseResultPackage( + artifact=ParseArtifact(add_dir=str(result_dir), dataframe=pd.DataFrame()), + chunks=[], + ), + generated_package=GeneratedResultPackage( + zip_file_path=str(zip_path), + checksum_value="abc", + statistics={}, + zip_size=3, + ), + job_id="job-m0", + result_storage_factory=FakeStorage, + job_context=SimpleNamespace( + job_metadata={ + "parse_track": "page_memory", + "ingestion_mode": "progressive", + } + ), # type: ignore[arg-type] + ) + assert persisted["enrichment_status"] == "text" + assert persisted["page_memory_phase_b"] == [ + PHASE_B_SOURCE_PDF_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, + ] + + +def test_page_memory_upload_without_manifest_does_not_mark_coarse( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + result_dir = tmp_path / "result" + result_dir.mkdir() + (result_dir / PHASE_B_SOURCE_PDF_REL).write_bytes(b"%PDF") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + + class FakeStorage: + def upload(self, **kwargs: object) -> SimpleNamespace: + del kwargs + return SimpleNamespace( + zip_key="results/job-1/result.zip", + raw_prefix="results/job-1/", + raw_files={"source.pdf": "results/job-1/source.pdf"}, + ) + + persisted: dict[str, object] = {} + + def capture_metadata(**kwargs: object) -> None: + persisted.update(kwargs["metadata_updates"]) # type: ignore[arg-type] + + monkeypatch.setattr( + success_finalization_mod, + "persist_job_metadata_updates", + capture_metadata, + ) + success_finalization_mod._upload_result_package( # noqa: SLF001 + result_package=ParseResultPackage( + artifact=ParseArtifact(add_dir=str(result_dir), dataframe=pd.DataFrame()), + chunks=[], + ), + generated_package=GeneratedResultPackage( + zip_file_path=str(zip_path), + checksum_value="abc", + statistics={}, + zip_size=3, + ), + job_id="job-1", + result_storage_factory=FakeStorage, + job_context=SimpleNamespace(job_metadata={"parse_track": "page_memory"}), # type: ignore[arg-type] + ) + assert persisted["result_raw_prefix"] == "results/job-1/" + assert "enrichment_status" not in persisted + assert "page_memory_phase_b" not in persisted + + +def test_result_storage_uploads_phase_b_internal_files(tmp_path: Path) -> None: + class FakeStorageAdapter: + def __init__(self) -> None: + self.uploaded_keys: list[str] = [] + + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: + del local_path, bucket + self.uploaded_keys.append(key) + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + result_dir = tmp_path / "result" + (result_dir / "phase_b").mkdir(parents=True) + (result_dir / "pages").mkdir() + (result_dir / PHASE_B_SOURCE_PDF_REL).write_bytes(b"source") + (result_dir / PHASE_B_MANIFEST_REL).write_text("{}", encoding="utf-8") + (result_dir / "pages" / "page-1.png").write_bytes(b"skip") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + + adapter = FakeStorageAdapter() + storage = JobResultStorage( + results_bucket="test-results", + storage_adapter=adapter, # type: ignore[arg-type] + ) + bundle = storage.upload( + job_id="job-1", + result_dir=str(result_dir), + zip_file_path=str(zip_path), + artifact_refs={ + PHASE_B_SOURCE_PDF_REL, + PHASE_B_MANIFEST_REL, + "pages/page-1.png", + }, + ) + + assert set(bundle.raw_files) == {PHASE_B_SOURCE_PDF_REL, PHASE_B_MANIFEST_REL} + assert "results/job-1/phase_b/manifest.json" in adapter.uploaded_keys + assert "results/job-1/pages/page-1.png" not in adapter.uploaded_keys + + +def test_scanned_progressive_m0_skips_ocr_and_builds_visual_page_chunks( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from app.services.document_agent.coordinator import ProfileCoordinator + from app.services.document_agent.manifest import DocumentProfile + from app.services.document_agent.registry import REGISTRY + from app.services.page_memory import memory_service, page_renderer + from app.services.page_memory.memory_service import PageMemoryInput + from shared.core.config import settings + from shared.services.ai import embedding as embedding_pkg + from shared.services.ai.embedding import DEFAULT_EMBEDDING_DIM + + ocr_called: list[bool] = [] + monkeypatch.setattr( + REGISTRY, + "dispatch", + lambda name, ctx, args: ( + ocr_called.append(True) + if name == "ocr.pages" + else SimpleNamespace(status="ok", payload={}, output_summary={}) + ), + ) + + source = tmp_path / "scanned.pdf" + source.write_bytes(b"%PDF") + coord = ProfileCoordinator( + pdf_path=str(source), + job_id="job-scan-m0", + output_dir=str(tmp_path / "_agent"), + db=None, + model="deepseek-flash", + settings={"skip_scanned_ocr": True}, + ) + coord.blackboard.page_count = 2 + coord.blackboard.document_profile = DocumentProfile( + is_scanned=True, + category=["04000000"], + category_label="Scanned Report", + routing_category="scanned", + language="en", + ) + coord._run_text_scan() # noqa: SLF001 + assert ocr_called == [] + assert coord.blackboard.page_full_text_cache == {} + + profile = _profile( + skeletons=[], + page_texts={}, + body_texts={}, + monkeypatch=monkeypatch, + ) + profile.is_scanned = True + profile.page_count = 2 + profile.blackboard = coord.blackboard + + captured_skip_ocr: list[object] = [] + + def fake_profile(*args: object, **kwargs: object) -> SimpleNamespace: + del args + captured_skip_ocr.append(kwargs.get("skip_scanned_ocr")) + return profile + + def fake_render(**kwargs: object) -> list[SimpleNamespace]: + out_dir = Path(str(kwargs["output_dir"])) / "pages" / "m0" + out_dir.mkdir(parents=True, exist_ok=True) + items: list[SimpleNamespace] = [] + for p in (1, 2): + img_path = out_dir / f"page-{p}.png" + img_path.write_bytes(b"\x89PNG\r\n\x1a\n" + bytes([p])) + items.append(SimpleNamespace(page_index=p, image_path=str(img_path))) + return items + + vec_p1 = [1.0] + [0.0] * (DEFAULT_EMBEDDING_DIM - 1) + vec_p2 = [0.0, 1.0] + [0.0] * (DEFAULT_EMBEDDING_DIM - 2) + + monkeypatch.setattr(settings, "VISUAL_RETRIEVAL_ENABLED", True) + monkeypatch.setattr( + memory_service, + "normalize_to_pdf", + lambda **kwargs: (str(source), "scanned.pdf"), + ) + monkeypatch.setattr(memory_service, "profile_document", fake_profile) + monkeypatch.setattr(page_renderer, "render_document_pages", fake_render) + monkeypatch.setattr( + embedding_pkg, + "embed_images", + lambda inputs, **kw: [vec_p1, vec_p2][: len(list(inputs))], + ) + + output_dir, dataframe = memory_service.run( + PageMemoryInput( + file_path=str(source), + filename="scanned.pdf", + output_dir=str(tmp_path / "out_scan"), + page_memory_config=PageMemoryConfig(asset_extraction_enabled=False), + ingestion_mode="progressive", + ) + ) + assert captured_skip_ocr == [True] + chunks = dataframe_to_chunks(dataframe) + assert len(chunks) == 2 + assert [c["content"] for c in chunks] == ["", ""] + assert [c["metadata"]["enrichment_state"] for c in chunks] == ["text", "text"] + assert chunks[0]["metadata"]["_page_embedding"][:2] == [1.0, 0.0] + assert chunks[1]["metadata"]["_page_embedding"][:2] == [0.0, 1.0] + assert ( + chunks[0]["metadata"]["page_assets"][0]["artifact_ref"] + == "page_citation_assets/page-1.png" + ) + assert ( + Path(output_dir) / "page_citation_assets" / "page-1.png" + ).is_file() + diff --git a/apps/worker/tests/contract/test_page_memory_root_prefix_contract.py b/apps/worker/tests/contract/test_page_memory_root_prefix_contract.py index 1c63b01d3..2c1385fbf 100644 --- a/apps/worker/tests/contract/test_page_memory_root_prefix_contract.py +++ b/apps/worker/tests/contract/test_page_memory_root_prefix_contract.py @@ -167,11 +167,14 @@ def retain_chapter_titles(**kwargs: object) -> list[PageTagResult]: image_files_map={}, table_files_map={}, ) - expected_prefix: list[int] = ( + expected_prefix_span: list[int] = ( list(range(1, chapter_pages[0])) if chapter_pages else list(range(1, page_count + 1)) ) + expected_prefix: list[int] = [ + page for page in expected_prefix_span if page not in set(toc_pages) + ] root_path: str = filename if chapter_pages else f"{filename}/Root" root_chunks = [chunk for chunk in published if chunk["path"] == root_path] assert len(root_chunks) == bool(expected_prefix) @@ -204,15 +207,13 @@ def retain_chapter_titles(**kwargs: object) -> list[PageTagResult]: covered_pages: set[int] = { page for chunk in published for page in chunk["metadata"]["page_nums"] } - assert covered_pages == set(range(1, page_count + 1)) - ( - set(toc_pages) - set(expected_prefix) - ) + assert covered_pages == set(range(1, page_count + 1)) - set(toc_pages) for chunk in published: for asset in chunk["metadata"]["page_assets"]: assert (tmp_path / asset["artifact_ref"]).is_file() -def test_staged_prefix_scope_retains_root_identity_and_printed_toc( +def test_staged_prefix_scope_retains_root_identity_and_excludes_toc_pages( tmp_path: Path, ) -> None: from app.services.page_memory._serialization import serialize_hierarchy_artifact @@ -240,12 +241,13 @@ def test_staged_prefix_scope_retains_root_identity_and_printed_toc( import sys from pathlib import Path sys.path.insert(0, 'apps/worker/scripts/page_memory') -from _debug_pm_shared import load_hierarchy_artifact, build_debug_coarse_scopes +from _debug_pm_shared import load_hierarchy_artifact +from app.services.page_memory._utils import build_hierarchy_scopes +from app.services.page_memory.phase_a import build_phase_a_manifest from debug_pm_stage3_scope_fine_hierarchy import _resolve_scope_processing_pages from debug_pm_stage4_assets import _load_scope_asset_context from app.services.page_memory._serialization import derive_hierarchy_page_scope, build_hierarchy_tree from app.services.page_memory.fine_hierarchy import compute_fat_leaf_pages -from types import SimpleNamespace path = Path(sys.argv[1]) _, skeletons = load_hierarchy_artifact(path) assert skeletons[0].level == 0 @@ -253,16 +255,25 @@ def test_staged_prefix_scope_retains_root_identity_and_printed_toc( assert derive_hierarchy_page_scope(skeletons=skeletons, page_count=407) == list(range(1, 20)) assert build_hierarchy_tree(skeletons) == {} assert compute_fat_leaf_pages(skeletons, min_pages=4, toc_pages=[17, 18]) == set() -anatomy = SimpleNamespace(toc_result=SimpleNamespace(toc_pages=[17, 18])) -scopes = build_debug_coarse_scopes(skeletons=skeletons, filename='spacex.pdf', page_count=407, anatomy=anatomy) +coarse_scopes = build_hierarchy_scopes(skeletons=skeletons, filename='spacex.pdf', page_count=407) +manifest = build_phase_a_manifest( + filename='spacex.pdf', + page_count=407, + skeletons=skeletons, + coarse_scopes=coarse_scopes, + next_title_by_path={}, + toc_pages=[17, 18], +) +scopes = manifest['scopes'] +expected_pages = [p for p in range(1, 20) if p not in {17, 18}] assert len(scopes) == 1 -assert scopes[0]['processing_pages'] == list(range(1, 20)) -assert scopes[0]['excluded_toc_pages'] == [] +assert scopes[0]['processing_pages'] == expected_pages +assert scopes[0]['excluded_toc_pages'] == [17, 18] pages, excluded = _resolve_scope_processing_pages(scope_meta={}, skeletons=skeletons, page_count=407, toc_pages=[17, 18]) -assert pages == list(range(1, 20)) -assert excluded == [] +assert pages == expected_pages +assert excluded == [17, 18] context = _load_scope_asset_context(scope_id='p1-19', scope_dir=path.parent, page_count=407, toc_pages=[17, 18]) -assert context.pages == list(range(1, 20)) +assert context.pages == expected_pages """, str(artifact_path), ], diff --git a/apps/worker/tests/contract/test_page_memory_scope_enrichment_task_contract.py b/apps/worker/tests/contract/test_page_memory_scope_enrichment_task_contract.py new file mode 100644 index 000000000..a1c6107b3 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_scope_enrichment_task_contract.py @@ -0,0 +1,1185 @@ +"""Phase B scope enrichment task contracts.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +from sqlalchemy import create_engine, select, update +from sqlalchemy.engine import make_url +from sqlalchemy.orm import Session + +from app.services.document_agent.manifest import ( + PageAnatomyMap, + PageFeature, + PageLabel, + ShardPlan, + TocResult, +) +from app.services.page_memory._utils import build_hierarchy_scopes +from app.services.page_memory.page_assets import PageAsset +from app.services.page_memory.page_renderer import PageRenderResult +from app.services.page_memory.page_tagger import PageTagResult +from app.services.page_memory.phase_a import ( + PHASE_B_ANATOMY_REL, + PHASE_B_BODY_TEXT_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_SOURCE_PDF_REL, + build_phase_a_manifest, +) +from app.services.page_memory.skeleton_extractor import SectionSkeleton +from shared.core.celery_app import celery_app +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentScopeEnrichment, +) +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.testing.contract_runtime import ( + get_contract_database_url, +) + + +class _InMemoryResultStorage: + def __init__(self, files: dict[str, bytes]) -> None: + self.files = dict(files) + self.uploaded_raw: dict[str, bytes] = {} + self.uploaded_zips: list[str] = [] + + def download_raw_to_temp( + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + suffix: str, + temp_dir: str, + ) -> str: + del job_id, raw_prefix + if relative_path not in self.files: + raise FileNotFoundError(relative_path) + target = Path(temp_dir) / f"_dl_{Path(relative_path).name}{suffix}" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(self.files[relative_path]) + return str(target) + + def upload_raw_file( + self, + *, + job_id: str, + relative_path: str, + local_file_path: str, + raw_prefix: str | None = None, + ) -> None: + del job_id, raw_prefix + payload = Path(local_file_path).read_bytes() + self.uploaded_raw[relative_path] = payload + self.files[relative_path] = payload + + def upload( + self, + *, + job_id: str, + result_dir: str, + zip_file_path: str, + artifact_refs: set[str] | None = None, + ) -> Any: + del result_dir, artifact_refs + from shared.services.storage.result_storage import UploadedResultBundle + + zip_key = f"results/{job_id}.zip" + self.uploaded_zips.append(zip_key) + Path(zip_file_path).unlink(missing_ok=True) + return UploadedResultBundle( + zip_key=zip_key, + raw_prefix=f"results/{job_id}/", + raw_files={}, + ) + + +def _seed_published_coarse_doc( + session: Session, + *, + suffix: str, + scope_states: list[tuple[str, int, int, str]], +) -> tuple[str, str, str]: + from sqlalchemy import text + + user_id = f"enrich-user-{suffix}" + job_id = f"enrich-job-{suffix}" + revision_id = f"enrich-rev-{suffix}" + session.execute( + text( + 'INSERT INTO "user" (id, name, email) ' + "VALUES (:id, :name, :email) ON CONFLICT (id) DO NOTHING" + ), + {"id": user_id, "name": "Enrich User", "email": f"{user_id}@test.local"}, + ) + job = Job( + job_id=job_id, + user_id=user_id, + job_type="document_ingestion", + status="done", + source_type="file", + webhook_enabled=False, + job_metadata={ + "namespace": f"ns-{suffix}", + "parse_track": "page_memory", + "source_file_name": "doc.pdf", + "result_raw_prefix": f"results/{job_id}/", + "enrichment_status": "coarse", + "page_memory_phase_b": [ + PHASE_B_SOURCE_PDF_REL, + PHASE_B_MANIFEST_REL, + PHASE_B_ANATOMY_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, + ], + "parsing_params": {"top_summary_use_llm": False}, + }, + ) + session.add(job) + session.flush() + job_result = JobResult( + id=revision_id, + job_id=job_id, + delivery_mode="url", + result_s3_key=f"results/{job_id}_coarse.zip", + result_size=10, + inline_payload={"checksum": "coarse"}, + ) + session.add(job_result) + session.flush() + + published = RetrievalPublicationService().publish_document_state( + session, + job_id=job_id, + job_result_id=revision_id, + chunks=[ + { + "chunk_id": f"coarse-{scope_id}", + "type": "page", + "content": f"raw {scope_id}", + "path": f"doc.pdf/{scope_id}", + "metadata": { + "scope_id": scope_id, + "enrichment_state": "coarse", + "page_nums": [page_start, page_end], + }, + } + for scope_id, page_start, page_end, _state in scope_states + ], + update_namespace_snapshot=False, + ) + assert published is not None and published.document_id is not None + for scope_id, page_start, page_end, state in scope_states: + session.add( + DocumentScopeEnrichment( + document_id=published.document_id, + job_result_id=revision_id, + scope_id=scope_id, + page_start=page_start, + page_end=page_end, + state=state, + ) + ) + session.commit() + return published.document_id, job_id, revision_id + + +def _build_storage_inputs() -> dict[str, bytes]: + import json + + skeletons = [ + SectionSkeleton( + section_path="doc.pdf/Parent", + title="Parent", + level=1, + start_page=1, + end_page=2, + parent_path="doc.pdf", + ), + SectionSkeleton( + section_path="doc.pdf/Parent/OnlyChild", + title="OnlyChild", + level=2, + start_page=1, + end_page=2, + parent_path="doc.pdf/Parent", + ), + SectionSkeleton( + section_path="doc.pdf/Parent/OnlyChild/GrandA", + title="GrandA", + level=3, + start_page=1, + end_page=2, + parent_path="doc.pdf/Parent/OnlyChild", + ), + SectionSkeleton( + section_path="doc.pdf/Parent/OnlyChild/GrandB", + title="GrandB", + level=3, + start_page=1, + end_page=2, + parent_path="doc.pdf/Parent/OnlyChild", + ), + SectionSkeleton( + section_path="doc.pdf/Other", + title="Other", + level=1, + start_page=3, + end_page=4, + parent_path="doc.pdf", + ), + ] + scopes = [ + build_hierarchy_scopes( + skeletons=skeletons[:4], + filename="doc.pdf", + page_count=4, + )[0], + build_hierarchy_scopes( + skeletons=skeletons[4:], + filename="doc.pdf", + page_count=4, + )[0], + ] + manifest = build_phase_a_manifest( + filename="doc.pdf", + page_count=4, + skeletons=skeletons, + coarse_scopes=scopes, + next_title_by_path={"doc.pdf/Parent/OnlyChild/GrandB": "Other"}, + toc_pages=[], + ) + anatomy = PageAnatomyMap( + job_id="j1", + file_path="source.pdf", + page_count=4, + page_features=[ + PageFeature( + page=page, + raw_text_length=20, + text_density=0.5, + image_coverage=0.1, + image_count=1, + table_count=0, + drawings_count=0, + orientation="portrait", + width=600.0, + height=800.0, + has_asset=(page == 1), + is_blank_like=False, + ) + for page in range(1, 5) + ], + page_labels=[PageLabel(page=page, kind="normal") for page in range(1, 5)], + toc_result=TocResult(toc_pages=[]), + shard_plan=ShardPlan(enabled=False, reason="not_needed"), + ) + return { + PHASE_B_SOURCE_PDF_REL: b"%PDF-1.4", + PHASE_B_MANIFEST_REL: json.dumps(manifest).encode("utf-8"), + PHASE_B_ANATOMY_REL: json.dumps(anatomy.to_dict()).encode("utf-8"), + PHASE_B_FULL_TEXT_REL: json.dumps( + {"1": "page one text", "2": "page two text", "3": "p3", "4": "p4"} + ).encode("utf-8"), + PHASE_B_BODY_TEXT_REL: json.dumps( + {"1": "body one", "2": "body two", "3": "b3", "4": "b4"} + ).encode("utf-8"), + } + + +def _install_fake_scope_pipeline(monkeypatch: pytest.MonkeyPatch) -> None: + from PIL import Image + + def fake_render( + *, + output_dir: str, + scope_id: str, + pages: list[int], + page_texts: dict[int, str], + **kwargs: object, + ) -> list[PageRenderResult]: + del kwargs + out: list[PageRenderResult] = [] + render_dir = Path(output_dir) / "pages" / scope_id + render_dir.mkdir(parents=True, exist_ok=True) + for page in pages: + img_path = render_dir / f"page-{page}.png" + Image.new("RGB", (40, 50), color="white").save(img_path) + out.append( + PageRenderResult( + page_index=page, + image_path=str(img_path), + raw_text=page_texts.get(page, ""), + width=600.0, + height=800.0, + is_landscape=False, + ) + ) + return out + + def fake_tag( + *, pages: list[PageRenderResult], **kwargs: object + ) -> list[PageTagResult]: + del kwargs + return [ + PageTagResult( + page_index=item.page_index, + summary=f"summary p{item.page_index}", + keywords=[f"kw{item.page_index}"], + strategy_used="fake", + ) + for item in pages + ] + + def fake_assets( + *, + rendered_pages: list[PageRenderResult], + output_dir: str, + **kwargs: object, + ) -> dict[int, list[PageAsset]]: + del kwargs + if not rendered_pages: + return {} + page = rendered_pages[0].page_index + img_dir = Path(output_dir) / "images" + img_dir.mkdir(parents=True, exist_ok=True) + img_file = img_dir / f"page-{page}-fig-1.png" + from PIL import Image as PilImage + + PilImage.new("RGB", (20, 20), color="blue").save(img_file) + return { + page: [ + PageAsset( + asset_id=f"asset_p{page}_1", + page_index=page, + asset_index=1, + kind="figure", + title="Figure 1", + summary="chart summary", + keywords=["chart"], + entities=[], + bbox_px=[0, 0, 20, 20], + width_px=40, + height_px=50, + width_pt=600.0, + height_pt=800.0, + image_uri=f"images/{img_file.name}", + ) + ] + } + + import app.services.page_memory.page_assets as page_assets_mod + import app.services.page_memory.page_renderer as page_renderer_mod + import app.services.page_memory.page_tagger as page_tagger_mod + + monkeypatch.setattr(page_renderer_mod, "render_document_pages", fake_render) + monkeypatch.setattr(page_tagger_mod, "tag_pages", fake_tag) + monkeypatch.setattr( + page_assets_mod, "extract_page_assets_from_renders", fake_assets + ) + + +def test_enrichment_task_is_routed_to_low_queue_and_does_not_fail_job() -> None: + from app.core.tasks.enrichment_tasks import ( + build_page_memory_structure, + enrich_page_memory_scope, + ) + + assert enrich_page_memory_scope.name == ( + "app.core.tasks.enrichment_tasks.enrich_page_memory_scope" + ) + assert build_page_memory_structure.name == ( + "app.core.tasks.enrichment_tasks.build_page_memory_structure" + ) + assert ( + celery_app.conf.task_routes[ + "app.core.tasks.enrichment_tasks.build_page_memory_structure" + ] + == {"queue": "document_ingestion_medium"} + ) + assert celery_app.conf.task_routes["app.core.tasks.enrichment_tasks.*"] == { + "queue": "document_ingestion_low" + } + assert enrich_page_memory_scope.finalize_job_on_failure is False + assert build_page_memory_structure.finalize_job_on_failure is False + assert ( + enrich_page_memory_scope._extract_job_id(("rev-1", "p1-2"), {}) is None # noqa: SLF001 + ) + + +def test_execute_scope_enrichment_claims_runs_subtree_and_completes_when_all_refined( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from app.services.page_memory.scope_enrichment_service import ( + execute_scope_enrichment, + ) + + del worker_contract_environment + _install_fake_scope_pipeline(monkeypatch) + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + storage = _InMemoryResultStorage(_build_storage_inputs()) + + try: + with Session(engine) as session: + doc_id, job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="ok", + scope_states=[("p1-2", 1, 2, "queued"), ("p3-4", 3, 4, "refined")], + ) + + result = execute_scope_enrichment( + job_result_id=revision_id, + scope_id="p1-2", + result_storage_factory=lambda: storage, + ) + assert result["status"] == "refined" + assert result["all_refined"] is True + assert "images/page-1-fig-1.png" in storage.uploaded_raw + assert "page_citation_assets/page-1.png" in storage.uploaded_raw + assert storage.uploaded_zips == [f"results/{job_id}.zip"] + + with Session(engine) as session: + db_chunks = list( + session.scalars( + select(DocumentChunk) + .where(DocumentChunk.document_id == doc_id) + .where(DocumentChunk.job_result_id == revision_id) + .order_by(DocumentChunk.sort_order, DocumentChunk.id) + ).all() + ) + p1_2_chunks = [ + c + for c in db_chunks + if (c.chunk_metadata or {}).get("scope_id") == "p1-2" + ] + p3_4_chunks = [ + c + for c in db_chunks + if (c.chunk_metadata or {}).get("scope_id") == "p3-4" + ] + page_chunks = [c for c in p1_2_chunks if c.chunk_type == "page"] + asset_chunks = [c for c in p1_2_chunks if c.chunk_type == "image"] + assert [c.source_chunk_path for c in page_chunks] == [ + "doc.pdf/Parent/GrandA", + "doc.pdf/Parent/GrandB", + ] + assert { + (c.chunk_metadata or {}).get("enrichment_state") + for c in page_chunks + } == {"refined"} + assert len(asset_chunks) == 1 + assert [c.chunk_id for c in p3_4_chunks] == ["coarse-p3-4"] + + row = session.execute( + select(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == revision_id) + .where(DocumentScopeEnrichment.scope_id == "p1-2") + ).scalar_one() + assert row.state == "refined" + assert row.attempts == 1 + assert row.refined_at is not None + + job = session.execute(select(Job).where(Job.job_id == job_id)).scalar_one() + assert (job.job_metadata or {}).get("enrichment_status") == "complete" + job_result = session.execute( + select(JobResult).where(JobResult.id == revision_id) + ).scalar_one() + assert job_result.result_s3_key == f"results/{job_id}.zip" + assert (job_result.result_size or 0) > 0 + assert (job_result.inline_payload or {}).get("checksum") not in ( + None, + "", + "coarse", + ) + + second = execute_scope_enrichment( + job_result_id=revision_id, + scope_id="p1-2", + result_storage_factory=lambda: storage, + ) + assert second["status"] == "skipped" + finally: + engine.dispose() + + +def test_execute_scope_enrichment_aborts_when_revision_is_superseded( + worker_contract_environment: None, +) -> None: + from app.services.page_memory.scope_enrichment_service import ( + execute_scope_enrichment, + ) + + del worker_contract_environment + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + try: + with Session(engine) as session: + doc_id, _job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="stale", + scope_states=[("p1-2", 1, 2, "queued")], + ) + newer_job = Job( + job_id="enrich-job-stale-newer", + user_id="enrich-user-stale", + job_type="document_ingestion", + status="done", + source_type="file", + webhook_enabled=False, + ) + session.add(newer_job) + session.flush() + newer = JobResult( + id="enrich-rev-stale-newer", + job_id=newer_job.job_id, + document_id=doc_id, + delivery_mode="url", + ) + session.add(newer) + session.flush() + doc = session.execute( + select(Document).where(Document.document_id == doc_id) + ).scalar_one() + doc.current_job_result_id = newer.id + session.commit() + + called = False + + def fail_if_called(*args: object, **kwargs: object) -> None: + nonlocal called + called = True + + result = execute_scope_enrichment( + job_result_id=revision_id, + scope_id="p1-2", + scope_publisher=fail_if_called, + ) + assert result["status"] == "aborted" + assert called is False + + with Session(engine) as session: + row = session.execute( + select(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == revision_id) + .where(DocumentScopeEnrichment.scope_id == "p1-2") + ).scalar_one() + assert row.state == "failed" + finally: + engine.dispose() + + +def test_execute_scope_enrichment_reclaims_stale_running_scope( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from datetime import timedelta + + from app.services.page_memory.scope_enrichment_service import ( + execute_scope_enrichment, + ) + from shared.core.config import app_config + from shared.utils.utc_now import utc_now_naive + + del worker_contract_environment + _install_fake_scope_pipeline(monkeypatch) + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + storage = _InMemoryResultStorage(_build_storage_inputs()) + published_calls: list[dict[str, Any]] = [] + + def fake_scope_publisher( + db: Session, + *, + scope: Any, + scope_id: str, + chunks: list[dict[str, Any]], + section_summaries: dict[str, str], + ) -> None: + del db, scope, chunks, section_summaries + published_calls.append({"scope_id": scope_id}) + + try: + with Session(engine) as session: + _doc_id, _job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="zombie", + scope_states=[("p1-2", 1, 2, "running")], + ) + # Simulate a crashed/OOM worker: the row is stuck in ``running`` with + # a claim watermark older than the Celery hard time limit, so no live + # task can still own it. + stale_at = utc_now_naive() - timedelta( + seconds=app_config.TASK_TIME_LIMIT_SECONDS + 60 + ) + session.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == revision_id) + .where(DocumentScopeEnrichment.scope_id == "p1-2") + .values(updated_at=stale_at) + ) + session.commit() + + result = execute_scope_enrichment( + job_result_id=revision_id, + scope_id="p1-2", + result_storage_factory=lambda: storage, + scope_publisher=fake_scope_publisher, + ) + assert result["status"] == "refined" + assert len(published_calls) == 1 + + with Session(engine) as session: + row = session.execute( + select(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == revision_id) + .where(DocumentScopeEnrichment.scope_id == "p1-2") + ).scalar_one() + assert row.state == "refined" + assert row.attempts == 1 + finally: + engine.dispose() + + +def test_execute_scope_enrichment_skips_fresh_running_scope( + worker_contract_environment: None, +) -> None: + from app.services.page_memory.scope_enrichment_service import ( + execute_scope_enrichment, + ) + + del worker_contract_environment + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + try: + with Session(engine) as session: + _doc_id, _job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="fresh-running", + scope_states=[("p1-2", 1, 2, "running")], + ) + + called = False + + def fail_if_called(*args: object, **kwargs: object) -> None: + nonlocal called + called = True + + result = execute_scope_enrichment( + job_result_id=revision_id, + scope_id="p1-2", + scope_publisher=fail_if_called, + ) + assert result["status"] == "skipped" + assert called is False + + with Session(engine) as session: + row = session.execute( + select(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == revision_id) + .where(DocumentScopeEnrichment.scope_id == "p1-2") + ).scalar_one() + assert row.state == "running" + finally: + engine.dispose() + + +def test_zip_regeneration_failure_does_not_mark_job_complete( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from app.services.page_memory import scope_enrichment_service as ses_mod + from app.services.page_memory.scope_enrichment_service import ( + execute_scope_enrichment, + ) + + del worker_contract_environment + _install_fake_scope_pipeline(monkeypatch) + monkeypatch.setattr( + ses_mod, + "_regenerate_and_upload_full_package", + lambda **_kw: (_ for _ in ()).throw(RuntimeError("zip boom")), + ) + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + storage = _InMemoryResultStorage(_build_storage_inputs()) + try: + with Session(engine) as session: + _doc_id, job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="zip-fail", + scope_states=[("p1-2", 1, 2, "queued")], + ) + result = execute_scope_enrichment( + job_result_id=revision_id, + scope_id="p1-2", + result_storage_factory=lambda: storage, + ) + assert result["status"] == "refined" + with Session(engine) as session: + job = session.execute(select(Job).where(Job.job_id == job_id)).scalar_one() + assert (job.job_metadata or {}).get("enrichment_status") == "coarse" + finally: + engine.dispose() + + +def test_enqueue_sends_scopes_low_priority_largest_first( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from shared.core import celery_app as celery_app_mod + from shared.services.jobs.lifecycle.scope_enrichments import ( + ENRICH_SCOPE_TASK_NAME, + LOW_ENRICHMENT_QUEUE, + enqueue_low_priority_scope_enrichments_after_commit, + ) + + del worker_contract_environment + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + sent: list[tuple[str, list[object], str | None]] = [] + + class _FakeCelery: + def send_task( + self, + name: str, + args: list[object] | None = None, + queue: str | None = None, + **kwargs: object, + ) -> None: + del kwargs + sent.append((name, list(args or []), queue)) + + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: _FakeCelery()) + try: + with Session(engine) as session: + _doc_id, _job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="enq-order", + scope_states=[ + ("p1-2", 1, 2, "coarse"), + ("p3-10", 3, 10, "coarse"), + ], + ) + + enqueue_low_priority_scope_enrichments_after_commit(revision_id) + + assert sent == [ + (ENRICH_SCOPE_TASK_NAME, [revision_id, "p3-10"], LOW_ENRICHMENT_QUEUE), + (ENRICH_SCOPE_TASK_NAME, [revision_id, "p1-2"], LOW_ENRICHMENT_QUEUE), + ] + with Session(engine) as session: + rows = { + r.scope_id: (r.state, r.priority) + for r in session.scalars( + select(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == revision_id + ) + ).all() + } + assert rows == {"p1-2": ("queued", "low"), "p3-10": ("queued", "low")} + finally: + engine.dispose() + + +def test_sweeper_reenqueues_only_revisions_stuck_at_text( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from datetime import timedelta + + from sqlalchemy import update + + from shared.core import celery_app as celery_app_mod + from shared.services.jobs.lifecycle.scope_enrichments import ( + BUILD_STRUCTURE_TASK_NAME, + sweep_progressive_enrichments_sync, + ) + from shared.utils.utc_now import utc_now_naive + + del worker_contract_environment + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + sent: list[tuple[str, list[object]]] = [] + + class _FakeCelery: + def send_task( + self, + name: str, + args: list[object] | None = None, + **kwargs: object, + ) -> None: + del kwargs + sent.append((name, list(args or []))) + + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: _FakeCelery()) + stale = utc_now_naive() - timedelta(hours=1) + try: + with Session(engine) as session: + _doc, stuck_job, stuck_rev = _seed_published_coarse_doc( + session, + suffix="sweep-text", + scope_states=[("p1-2", 1, 2, "refined")], + ) + _doc, done_job, done_rev = _seed_published_coarse_doc( + session, + suffix="sweep-done", + scope_states=[("p1-2", 1, 2, "refined")], + ) + for job_id, status in ((stuck_job, "text"), (done_job, "complete")): + job = session.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one() + job.job_metadata = { + **dict(job.job_metadata or {}), + "enrichment_status": status, + } + session.flush() + session.execute( + update(Job) + .where(Job.job_id.in_([stuck_job, done_job])) + .values(updated_at=stale) + ) + session.commit() + + sweep_progressive_enrichments_sync() + + structure_sends = [args for name, args in sent if name == BUILD_STRUCTURE_TASK_NAME] + assert [stuck_rev] in structure_sends + assert [done_rev] not in structure_sends + finally: + engine.dispose() + + +def test_enqueue_failure_reverts_unsent_scopes_to_coarse_and_sweeper_retries_failed_scopes( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from shared.core import celery_app as celery_app_mod + from shared.services.jobs.lifecycle.scope_enrichments import ( + enqueue_low_priority_scope_enrichments_after_commit, + sweep_progressive_enrichments_sync, + ) + + del worker_contract_environment + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + + class _FailingCelery: + def __init__(self) -> None: + self.calls = 0 + + def send_task(self, *args: object, **kwargs: object) -> None: + del args, kwargs + self.calls += 1 + if self.calls == 2: + raise RuntimeError("broker down") + + failing = _FailingCelery() + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: failing) + + try: + with Session(engine) as session: + _doc_id, _job_id, revision_id = _seed_published_coarse_doc( + session, + suffix="enq-revert", + scope_states=[ + ("p1-4", 1, 4, "coarse"), + ("p5-6", 5, 6, "coarse"), + ], + ) + + sent = enqueue_low_priority_scope_enrichments_after_commit(revision_id) + assert sent == ["p1-4"] + + with Session(engine) as session: + rows = { + r.scope_id: r + for r in session.scalars( + select(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == revision_id + ) + ).all() + } + assert rows["p1-4"].state == "queued" + # Unsent p5-6 reverted to coarse instead of stuck in queued! + assert rows["p5-6"].state == "coarse" + assert rows["p5-6"].priority is None + # Mark p1-4 as failed with attempts=1 (< MAX_SCOPE_ENRICHMENT_ATTEMPTS=3). + rows["p1-4"].state = "failed" + rows["p1-4"].attempts = 1 + session.commit() + + sent_by_sweep: list[tuple[str, str]] = [] + + class _OkCelery: + def send_task( + self, + name: str, + args: list[object] | None = None, + **kwargs: object, + ) -> None: + del name, kwargs + if args and len(args) == 2: + sent_by_sweep.append((str(args[0]), str(args[1]))) + + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: _OkCelery()) + sweep_res = sweep_progressive_enrichments_sync() + assert sweep_res["status"] == "success" + retried_for_rev = [ + scope_id for rev, scope_id in sent_by_sweep if rev == revision_id + ] + assert set(retried_for_rev) == {"p1-4", "p5-6"} + finally: + engine.dispose() + + +def test_execute_structure_enrichment_upgrades_m0_text_to_coarse_and_enqueues_scopes( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import json + + from app.services.page_memory import scope_enrichment_service as ses_mod + from app.services.page_memory.phase_a import PHASE_B_BLACKBOARD_REL + from app.services.page_memory.scope_enrichment_service import ( + execute_structure_enrichment, + ) + from shared.core import celery_app as celery_app_mod + from shared.services.jobs.lifecycle.scope_enrichments import ( + sweep_progressive_enrichments_sync, + ) + + del worker_contract_environment + + database_url = make_url(get_contract_database_url()).set( + drivername="postgresql+psycopg2" + ) + engine = create_engine(database_url) + sent_tasks: list[tuple[str, list[object]]] = [] + + class _FakeCelery: + def send_task( + self, + name: str, + args: list[object] | None = None, + **kwargs: object, + ) -> None: + del kwargs + sent_tasks.append((name, list(args or []))) + + monkeypatch.setattr(celery_app_mod, "get_celery_app", lambda: _FakeCelery()) + + skeletons = [ + SectionSkeleton("doc.pdf/Ch1", 1, 1, 2, "Ch1", "doc.pdf"), + SectionSkeleton("doc.pdf/Ch2", 1, 3, 4, "Ch2", "doc.pdf"), + ] + anatomy = PageAnatomyMap( + job_id="j-struct", + file_path="source.pdf", + page_count=4, + page_features=[], + page_labels=[], + toc_result=TocResult(toc_pages=[]), + shard_plan=ShardPlan(enabled=False, reason="not_needed"), + ) + from types import SimpleNamespace + + fake_profile = SimpleNamespace( + page_count=4, + anatomy=anatomy, + toc=SimpleNamespace(has_toc=True), + page_full_text_cache={1: "p1 text", 2: "p2 text", 3: "p3 text", 4: "p4 text"}, + page_body_text_cache={1: "p1 body", 2: "p2 body", 3: "p3 body", 4: "p4 body"}, + ) + monkeypatch.setattr( + ses_mod, + "_resume_profile_from_m0_workspace", + lambda **_kw: fake_profile, + ) + import app.services.page_memory.skeleton_extractor as skel_mod + + monkeypatch.setattr( + skel_mod, "extract_section_skeletons", lambda **_kw: skeletons + ) + + storage = _InMemoryResultStorage( + { + PHASE_B_SOURCE_PDF_REL: b"%PDF-1.4", + PHASE_B_BLACKBOARD_REL: json.dumps({"page_count": 4}).encode("utf-8"), + PHASE_B_FULL_TEXT_REL: json.dumps( + {"1": "p1 text", "2": "p2 text", "3": "p3 text", "4": "p4 text"} + ).encode("utf-8"), + PHASE_B_BODY_TEXT_REL: json.dumps( + {"1": "p1 body", "2": "p2 body", "3": "p3 body", "4": "p4 body"} + ).encode("utf-8"), + } + ) + + try: + from sqlalchemy import text + + user_id = "struct-user" + job_id = "struct-job" + revision_id = "struct-rev" + with Session(engine) as session: + session.execute( + text( + 'INSERT INTO "user" (id, name, email) ' + "VALUES (:id, :name, :email) ON CONFLICT (id) DO NOTHING" + ), + {"id": user_id, "name": "Struct", "email": "struct@test.local"}, + ) + job = Job( + job_id=job_id, + user_id=user_id, + job_type="document_ingestion", + status="done", + source_type="file", + webhook_enabled=False, + job_metadata={ + "namespace": "ns-struct", + "parse_track": "page_memory", + "ingestion_mode": "progressive", + "source_file_name": "doc.pdf", + "result_raw_prefix": f"results/{job_id}/", + "enrichment_status": "text", + "page_memory_phase_b": [ + PHASE_B_SOURCE_PDF_REL, + PHASE_B_BLACKBOARD_REL, + PHASE_B_FULL_TEXT_REL, + PHASE_B_BODY_TEXT_REL, + ], + }, + ) + session.add(job) + session.flush() + jr = JobResult( + id=revision_id, + job_id=job_id, + delivery_mode="url", + document_metadata={"enrichment_status": "text"}, + result_s3_key=f"results/{job_id}_m0.zip", + result_size=10, + inline_payload={"checksum": "m0"}, + ) + session.add(jr) + session.flush() + published = RetrievalPublicationService().publish_document_state( + session, + job_id=job_id, + job_result_id=revision_id, + chunks=[ + { + "chunk_id": f"m0-p{p}", + "type": "page", + "content": f"p{p} text", + "path": f"doc.pdf/p{p}", + "metadata": { + "scope_id": f"p{p}-{p}", + "enrichment_state": "text", + "page_nums": [p], + }, + } + for p in range(1, 5) + ], + update_namespace_snapshot=False, + ) + assert published is not None and published.document_id is not None + doc_id = published.document_id + session.commit() + + result = execute_structure_enrichment( + job_result_id=revision_id, + result_storage_factory=lambda: storage, + ) + assert result["status"] == "coarse" + assert result["scope_count"] == 2 + assert PHASE_B_MANIFEST_REL in storage.uploaded_raw + assert PHASE_B_ANATOMY_REL in storage.uploaded_raw + assert len(sent_tasks) == 2 + + with Session(engine) as session: + chunks_after = list( + session.scalars( + select(DocumentChunk) + .where(DocumentChunk.document_id == doc_id) + .where(DocumentChunk.job_result_id == revision_id) + .order_by(DocumentChunk.sort_order) + ).all() + ) + assert [c.source_chunk_path for c in chunks_after] == [ + "doc.pdf/Ch1", + "doc.pdf/Ch2", + ] + assert { + (c.chunk_metadata or {}).get("enrichment_state") + for c in chunks_after + } == {"coarse"} + scope_rows = list( + session.scalars( + select(DocumentScopeEnrichment) + .where( + DocumentScopeEnrichment.job_result_id == revision_id + ) + .order_by(DocumentScopeEnrichment.scope_id) + ).all() + ) + assert [(r.scope_id, r.state) for r in scope_rows] == [ + ("p1-2", "queued"), + ("p3-4", "queued"), + ] + job = session.execute( + select(Job).where(Job.job_id == job_id) + ).scalar_one() + assert (job.job_metadata or {}).get("enrichment_status") == "coarse" + + # Reset to text while structure_claimed_at is still fresh: both a duplicate + # structure task and sweep must skip it rather than re-running or re-enqueuing. + with Session(engine) as session: + job = session.execute(select(Job).where(Job.job_id == job_id)).scalar_one() + jr = session.execute(select(JobResult).where(JobResult.id == revision_id)).scalar_one() + job.job_metadata = {**dict(job.job_metadata or {}), "enrichment_status": "text"} + jr.document_metadata = {**dict(jr.document_metadata or {}), "enrichment_status": "text"} + session.commit() + + dup_res = execute_structure_enrichment( + job_result_id=revision_id, + result_storage_factory=lambda: storage, + ) + assert dup_res["status"] == "skipped" + sweep_while_claimed = sweep_progressive_enrichments_sync(structure_timeout_seconds=0) + assert revision_id not in sweep_while_claimed["reenqueued_structures"] + finally: + engine.dispose() + + diff --git a/apps/worker/tests/contract/test_toc_graft_contract.py b/apps/worker/tests/contract/test_toc_graft_contract.py index 4777c5b9a..c31e9e478 100644 --- a/apps/worker/tests/contract/test_toc_graft_contract.py +++ b/apps/worker/tests/contract/test_toc_graft_contract.py @@ -622,7 +622,7 @@ def test_page_parallel_still_cuts_primary_window() -> None: page_texts=page_texts, ) ch2 = next(item for item in skeletons if item.title == "Ch2") - assert ch2.end_page == 199 + assert ch2.end_page == 209 assert any(item.title == "App" for item in skeletons) diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index 9e66fdeac..247216b0f 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -6,7 +6,7 @@ The templates intentionally omit `S3_ACCESS_KEY_ID` and `S3_SECRET_ACCESS_KEY`. The environment-specific Secrets Manager secret supplied to the renderer must be a JSON secret with these keys: -- API: `DATABASE_URL`, `REDIS_HOST`, `REDIS_PORT`, `REDIS_PASSWORD`, `CELERY_REDIS_URL`, `SECRET_KEY`, `DS_KEY`, `ALI_API_KEYS`, `ARK_API_KEY`, `GPT_API_KEY`, `CURSOR_API_KEY`, `MINERU_API_KEYS`, `STRIPE_SECRET_KEY`, `STRIPE_WEBHOOK_SECRET`, `WEBHOOK_MASTER_KEY`, `LOGFIRE_TOKEN`, `QSTASH_TOKEN`, `QSTASH_CURRENT_SIGNING_KEY`, `QSTASH_NEXT_SIGNING_KEY` +- API: `DATABASE_URL`, `REDIS_HOST`, `REDIS_PORT`, `REDIS_PASSWORD`, `CELERY_REDIS_URL`, `SECRET_KEY`, `DS_KEY`, `ALI_API_KEYS`, `ARK_API_KEY`, `GPT_API_KEY`, `CURSOR_API_KEY`, `VERCEL_KEY`, `MINERU_API_KEYS`, `STRIPE_SECRET_KEY`, `STRIPE_WEBHOOK_SECRET`, `WEBHOOK_MASTER_KEY`, `LOGFIRE_TOKEN`, `QSTASH_TOKEN`, `QSTASH_CURRENT_SIGNING_KEY`, `QSTASH_NEXT_SIGNING_KEY` - Worker: the API keys above plus `CELERY_REDIS_PASSWORD` and `ILOVEAPI_KEYS` Render only after the secret and log groups exist, substituting the exact immutable ECR image digests and IAM role/secret ARNs: diff --git a/deploy/ecs/task-definition-api.staging.json b/deploy/ecs/task-definition-api.staging.json index 0e6cf9a97..183856c61 100644 --- a/deploy/ecs/task-definition-api.staging.json +++ b/deploy/ecs/task-definition-api.staging.json @@ -44,7 +44,7 @@ {"name": "DS_URL", "value": "https://api.deepseek.com/v1/chat/completions"}, {"name": "ALI_URL", "value": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1/chat/completions"}, {"name": "ARK_URL", "value": "https://ark.cn-beijing.volces.com/api/v3/chat/completions"}, - {"name": "EMBEDDING_MODEL", "value": "text-embedding-v4"}, + {"name": "EMBEDDING_MODEL", "value": "google/gemini-embedding-2"}, {"name": "RATE_LIMIT_ENABLED", "value": "true"}, {"name": "TELEMETRY_ENABLED", "value": "false"}, {"name": "API_STANDALONE_MODE_ENABLED", "value": "true"}, @@ -77,7 +77,8 @@ {"name": "QSTASH_TOKEN", "valueFrom": "${SECRETS_ARN}:QSTASH_TOKEN::"}, {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"}, - {"name": "CURSOR_API_KEY", "valueFrom": "${SECRETS_ARN}:CURSOR_API_KEY::"} + {"name": "CURSOR_API_KEY", "valueFrom": "${SECRETS_ARN}:CURSOR_API_KEY::"}, + {"name": "VERCEL_KEY", "valueFrom": "${SECRETS_ARN}:VERCEL_KEY::"} ], "healthCheck": { "command": ["CMD-SHELL", "curl -f http://localhost:5005/health || exit 1"], diff --git a/deploy/ecs/task-definition-worker.staging.json b/deploy/ecs/task-definition-worker.staging.json index bebb8f415..f893828d0 100644 --- a/deploy/ecs/task-definition-worker.staging.json +++ b/deploy/ecs/task-definition-worker.staging.json @@ -56,7 +56,7 @@ {"name": "DS_URL", "value": "https://api.deepseek.com/v1/chat/completions"}, {"name": "ALI_URL", "value": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1/chat/completions"}, {"name": "ARK_URL", "value": "https://ark.cn-beijing.volces.com/api/v3/chat/completions"}, - {"name": "EMBEDDING_MODEL", "value": "text-embedding-v4"}, + {"name": "EMBEDDING_MODEL", "value": "google/gemini-embedding-2"}, {"name": "HIERARCHY_LLM_MODEL", "value": "deepseek-chat"}, {"name": "NORMOL_MODEL", "value": "deepseek-chat"}, {"name": "IMAGE_MODEL", "value": "deepseek-flash"}, @@ -82,7 +82,8 @@ {"name": "QSTASH_TOKEN", "valueFrom": "${SECRETS_ARN}:QSTASH_TOKEN::"}, {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"}, - {"name": "CURSOR_API_KEY", "valueFrom": "${SECRETS_ARN}:CURSOR_API_KEY::"} + {"name": "CURSOR_API_KEY", "valueFrom": "${SECRETS_ARN}:CURSOR_API_KEY::"}, + {"name": "VERCEL_KEY", "valueFrom": "${SECRETS_ARN}:VERCEL_KEY::"} ], "healthCheck": { "command": ["CMD-SHELL", "python -c \"from shared.services.worker_health import assert_worker_healthy; assert_worker_healthy()\""], diff --git a/deploy/ecs/test_render_task_definitions.py b/deploy/ecs/test_render_task_definitions.py index af497c582..181470e1b 100644 --- a/deploy/ecs/test_render_task_definitions.py +++ b/deploy/ecs/test_render_task_definitions.py @@ -44,7 +44,7 @@ "ARK_URL": "https://ark.cn-beijing.volces.com/api/v3/chat/completions", "BROKER_POOL_LIMIT": "5", "DS_URL": "https://api.deepseek.com/v1/chat/completions", - "EMBEDDING_MODEL": "text-embedding-v4", + "EMBEDDING_MODEL": "google/gemini-embedding-2", "S3_TEMP_PATH": "/tmp", "SUPPORTED_EXTENSIONS": ".doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.jpg,.jpeg,.png,.pptx,.md", "TMP_PATH": "/tmp/aismart_bid", diff --git a/deprecated/mapnav/README.md b/deprecated/mapnav/README.md deleted file mode 100644 index a7d0c19bf..000000000 --- a/deprecated/mapnav/README.md +++ /dev/null @@ -1,26 +0,0 @@ -# Archived map-nav retrieval route - -This directory holds the retired checklist map-nav episode (PLANNER / HARVEST / CONTROL). - -It is **not imported** by production code. Lint, typecheck, and pytest skip it. - -## Why it was archived - -Retrieval now has two live routes: - -- `use_agentic is False` → classic map-unit BM25 -- otherwise → `agent_explore` (default harness: `cursor_sdk`) - -The `RETRIEVAL_AGENTIC_ROUTER=mapnav` switch was removed. Setting that env var has no effect. - -Shared scoring used by publication and classic recall was extracted first, to `packages/shared-python/shared/services/retrieval/scoring/`. - -## How to restore (manual) - -1. Copy these files back to their original paths: - - `nav/` → `packages/shared-python/shared/services/retrieval/nav/` - - `nav_config.py`, `nav_snapshot.py`, `nav_bridge.py`, `nav_llm_backend.py` → `packages/shared-python/shared/services/retrieval/` - - `trace_mapnav.py` → `packages/shared-python/shared/services/retrieval/trace/mapnav.py` -2. Reverse the Phase 1 import moves: leftover map-nav modules expect `shared.services.retrieval.nav.*` for types that now live under `scoring/`. -3. Re-attach a third branch in `execution/routes.py` (`use_agentic is False` → classic, else map-nav or `agent_explore`). -4. Restore the archived tests from `tests/` and stop excluding this directory from pytest. diff --git a/deprecated/mapnav/nav/__init__.py b/deprecated/mapnav/nav/__init__.py deleted file mode 100644 index 50aa1b62d..000000000 --- a/deprecated/mapnav/nav/__init__.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Recursive-dispatch map navigation for RealData experiments. - -LEGACY, PENDING REPLACEMENT: this package implements the map-nav -PLANNER/HARVEST/CONTROL episode -(``run_nav_episode``), the current default agentic retrieval route via -``execution.routes._run_mapnav_route``. It will be superseded by -``agent_explore`` after its evaluation gate, then trimmed in -Phase 5 to whatever this package still owns and nothing else uses (the BM25 -scorer in ``knowhere_hybrid.py``, the ``NodeFilter`` predicate in -``nav_node_filter.py``, and snapshot loading are already planned to be -reused by the new ``agent_tools/``, not deleted). Do not add new PLANNER / -HARVEST / CONTROL capabilities here; new agentic-retrieval work belongs in -``shared/services/retrieval/agent_tools/`` and -``shared/services/retrieval/agent_explore/``. -""" - -from .nav_types import ( - ActionKind, - NavConfig, - NavState, - SubgoalResult, -) -from .nav_agent import run_nav_episode -from .nav_plan import ( - Contract, - RetrievalPlan, - Subgoal, - bind_slots, - plan_query, -) -from .nav_orchestrate import execute_plan - -__all__ = [ - "ActionKind", - "NavConfig", - "NavState", - "SubgoalResult", - "run_nav_episode", - "Contract", - "RetrievalPlan", - "Subgoal", - "bind_slots", - "plan_query", - "execute_plan", -] diff --git a/deprecated/mapnav/nav/_compat.py b/deprecated/mapnav/nav/_compat.py deleted file mode 100644 index 7d6e31dba..000000000 --- a/deprecated/mapnav/nav/_compat.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Minimal stand-ins for experiment-repo ``agent_delivery`` symbols. - -Production map-nav runs with ``toolspace=ProviderToolSpace(...)`` and -``compose_answer=False``, so most of these are type/shape only. Query-time -map scoring uses the persisted path+content BM25 corpus only. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple - - -@dataclass -class Chunk: - node_id: str - doc_id: str - text: str - line_ids: Tuple[int, ...] - section_id: Optional[str] = None - text_line_id_groups: Optional[Tuple[Tuple[int, ...], ...]] = None - - -@dataclass -class AgentStep: - step_idx: int - action: str - detail: Dict[str, object] = field(default_factory=dict) - - -@dataclass -class EpisodeResult: - representation: str - steps: List[AgentStep] - scored_chunks: List[Tuple[Chunk, float]] - kept_chunks: List[Chunk] - evidence_text: str - evidence_chars_actual: int - retrieved_nodes: List[str] - composed_answer: str = "" - section_ids: List[str] = field(default_factory=list) - trajectory_length: int = 0 - truncated_last: bool = False - phase_timings: Dict[str, float] = field(default_factory=dict) - stop_reason: str = "completed" - - -# Duck-typed ToolSpace is provided by ``ProviderToolSpace``; keep a type alias. -ToolSpace = Any -HierarchicalTools = Any - - -def line_node_id(doc_id: str, line_id: int) -> str: - """Experiment-corpus helper; ProviderToolSpace never hits this path.""" - raise NotImplementedError( - "line_node_id is not available in Knowhere nav; use ProviderToolSpace" - ) - - -def compose_answer_llm(*_args: Any, **_kwargs: Any) -> str: - raise NotImplementedError( - "compose_answer_llm is not wired in Knowhere retrieval; pass compose_answer=False" - ) - - -def snapshot_usage() -> Dict[str, Any]: - """Fallback when no ``nav_token_episode`` is active.""" - return {} - - -def record_usage(*_args: Any, **_kwargs: Any) -> None: - return None - - -def load_llm_env() -> None: - return None - - -def require_llm_env(*_args: Any, **_kwargs: Any) -> None: - return None - - -def make_openai_client(*_args: Any, **_kwargs: Any) -> Any: - raise NotImplementedError("Install set_nav_chat_backend(...) for Knowhere nav LLM") - - -def resolve_thinking_mode(raw: Optional[str]) -> str: - text = str(raw or "").strip().lower() - if text in {"1", "true", "yes", "on", "enabled", "enable"}: - return "enabled" - if text in {"0", "false", "no", "off", "disabled", "disable"}: - return "disabled" - return "disabled" - - -def chat_thinking_extra(*, mode: str, model: str = "") -> Dict[str, Any]: - del model - typ = "enabled" if mode == "enabled" else "disabled" - return {"extra_body": {"thinking": {"type": typ}}} - - -def resolve_chat_credentials(*_args: Any, **_kwargs: Any) -> Tuple[str, str]: - return "", "" - - -def cached_chat_completion(*_args: Any, **_kwargs: Any) -> Dict[str, Any]: - raise NotImplementedError( - "Install set_nav_chat_backend(...) before calling nav_chat in Knowhere" - ) diff --git a/deprecated/mapnav/nav/nav_actions.py b/deprecated/mapnav/nav/nav_actions.py deleted file mode 100644 index 8af8b07d9..000000000 --- a/deprecated/mapnav/nav/nav_actions.py +++ /dev/null @@ -1,215 +0,0 @@ -from __future__ import annotations - -import os -from typing import Any, List, Optional - -from .nav_address import is_dispatch_only_node -from .nav_knowhere import is_root_section -from .nav_projection import format_harvested_tag, format_hit_tag -from .nav_types import ( - ActionKind, - LegalAction, - NavConfig, - NavState, - Projection, - SectionView, -) - - -def _env_enabled(name: str, default: str = "1") -> bool: - return os.environ.get(name, default).strip().lower() not in {"0", "false", "no", "off"} - - -def build_legal_actions( - state: NavState, - projection: Projection, - *, - config: NavConfig, - depth: int = 0, - ts: Any = None, -) -> List[LegalAction]: - """Every visible node is actionable: COLLECT + DISPATCH (when allowed) + FINISH. - - No action-space top-K: visibility is governed only by map display budget folding. - FINISH is always available; the LLM decides when to exit the current scope. - DISPATCH never targets the current scope root (no self-dispatch loop). - Document / namespace nodes are DISPATCH-only (level registry via ``ts``). - """ - actions: List[LegalAction] = [] - filter_collected = _env_enabled("NAV_FILTER_COLLECTED_SECTIONS") - collected_sids = set(state.collected_section_ids) | { - str(h.get("section_id") or "").strip() - for h in state.action_history - if h.get("kind") == "collect" and int(h.get("n_added", 0) or 0) > 0 - } - blocked = set(state.blocked_collect_section_ids) - map_scores = dict(getattr(state, "map_scores", {}) or {}) - unit_scores = dict(getattr(state, "unit_scores", {}) or {}) - highlight_set = set(projection.highlight_ids or state.highlight_ids or []) - scope_id = str(state.current_scope or projection.scope or "").strip() - - # DISPATCH only at depth 0 unless recursive dispatch is enabled and under max depth. - can_dispatch = depth == 0 or ( - bool(config.enable_recursive_dispatch) and depth < int(config.max_dispatch_depth) - ) - - rows = list(projection.tree_sections) or list(projection.visible_sections) - collect_i = 1 - dispatch_i = 1 - - def view_score(view: SectionView) -> float: - score = float(view.score) - if score <= 0.0: - score = float(map_scores.get(view.section_id, 0.0) or 0.0) - if score <= 0.0: - score = float(unit_scores.get(view.section_id, 0.0) or 0.0) - self_key = f"{view.section_id}__self" - if self_key in unit_scores: - score = max(score, float(unit_scores[self_key] or 0.0)) - return score - - for view in rows: - sid = view.section_id - label = view.title or view.preview or sid - score = view_score(view) - is_hit = sid in highlight_set or bool(view.is_highlight) - - collect_blocked = filter_collected and sid in ( - collected_sids | blocked - ) - if is_dispatch_only_node(ts, sid) or is_root_section(ts, sid): - collect_blocked = True - if not collect_blocked: - actions.append( - LegalAction( - action_id=f"C{collect_i}", - kind=ActionKind.COLLECT, - section_id=sid, - label=label, - score=score, - metadata={ - "map_id": view.map_id, - "highlight": is_hit, - "multi": True, - }, - ) - ) - collect_i += 1 - - if ( - can_dispatch - and view.has_children - and sid not in collected_sids - and sid != scope_id - ): - actions.append( - LegalAction( - action_id=f"D{dispatch_i}", - kind=ActionKind.DISPATCH, - section_id=sid, - label=label, - score=score, - metadata={ - "map_id": view.map_id, - "highlight": is_hit, - "multi": True, - }, - ) - ) - dispatch_i += 1 - - actions.append( - LegalAction( - action_id="F1", - kind=ActionKind.FINISH, - label="finish navigation and pack final evidence budget", - ) - ) - return actions - - -def format_actionable_map_observation( - projection: Projection, - actions: List[LegalAction], - *, - inline_summary: bool = False, -) -> str: - """Tree observation: each node line carries its legal action IDs.""" - by_sid: dict[str, list[LegalAction]] = {} - global_actions: List[LegalAction] = [] - for act in actions: - if act.kind == ActionKind.FINISH or not act.section_id: - global_actions.append(act) - continue - by_sid.setdefault(str(act.section_id), []).append(act) - - kind_label = { - ActionKind.COLLECT: "collect", - ActionKind.DISPATCH: "dispatch", - } - - def node_actions(sid: str) -> str: - acts = by_sid.get(sid) or [] - if not acts: - return "none" - order = {ActionKind.COLLECT: 0, ActionKind.DISPATCH: 1} - acts = sorted(acts, key=lambda a: (order.get(a.kind, 9), a.action_id)) - return ", ".join( - f"{kind_label.get(a.kind, a.kind.value)}={a.action_id}" for a in acts - ) - - rows = list(projection.tree_sections) or list(projection.visible_sections) - lines = [ - f"doc_id={projection.doc_id}", - f"scope={projection.scope or ''}", - "Each visible section appears once. Choose action IDs attached to the relevant line.", - ] - for view in rows: - indent = " " * max(0, int(view.depth_from_scope)) - leaf_tag = " [Leaf]" if not view.has_children else "" - hit_tag = format_hit_tag(is_highlight=bool(view.is_highlight)) - harvested_tag = format_harvested_tag(getattr(view, "harvested_by", "") or "") - map_id = view.map_id or "?" - lines.append( - f"{indent}[{map_id}] {view.title or view.section_id}" - f"{leaf_tag}{hit_tag}{harvested_tag} actions: {node_actions(view.section_id)}" - ) - if inline_summary and view.summary: - lines.append(f"{indent} summary: {view.summary}") - - if global_actions: - lines.append("") - lines.append("Global actions:") - for act in global_actions: - if act.kind == ActionKind.FINISH: - lines.append(f" finish={act.action_id}") - else: - lines.append(f" {act.action_id} ({act.kind.value})") - # SEARCH_* is harvest ``search_assets`` (nav_assets), not a map-row action. - return "\n".join(lines) - - -def action_by_id(actions: List[LegalAction], action_id: Optional[str]) -> Optional[LegalAction]: - aid = (action_id or "").strip().upper() - if not aid: - return None - for a in actions: - if a.action_id.upper() == aid: - return a - return None - - -def actions_by_ids(actions: List[LegalAction], ids: List[str]) -> List[LegalAction]: - """Resolve a multi-id selection (COLLECT / DISPATCH), preserving request order.""" - by_id = {a.action_id.upper(): a for a in actions} - out: List[LegalAction] = [] - seen: set[str] = set() - for raw in ids: - aid = (raw or "").strip().upper() - if not aid or aid in seen: - continue - act = by_id.get(aid) - if act is not None: - out.append(act) - seen.add(aid) - return out diff --git a/deprecated/mapnav/nav/nav_address.py b/deprecated/mapnav/nav/nav_address.py deleted file mode 100644 index 0e13d8b46..000000000 --- a/deprecated/mapnav/nav/nav_address.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Production address model for MAP-NAV. - -Levels match Knowhere's real keys (see ``document.py``): - -- ``NAMESPACE`` — ``(user_id, namespace)`` scope. Not a row; represented as - empty ``scope`` / empty node id, never a synthetic ``__corpus__:__root``. -- ``DOCUMENT`` — ``Document.document_id`` -- ``SECTION`` — ``DocumentSection.section_id`` -- ``CHUNK`` — ``DocumentChunk.chunk_id`` - -Hierarchy comes from the provider's level registry and ``parent_section_id``, -not from parsing magic suffixes out of a string id. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from enum import Enum -from typing import Any, Optional - - -class NavLevel(str, Enum): - NAMESPACE = "namespace" - DOCUMENT = "document" - SECTION = "section" - CHUNK = "chunk" - - -@dataclass(frozen=True) -class NavAddress: - level: NavLevel - id: str = "" - - def __post_init__(self) -> None: - object.__setattr__(self, "id", str(self.id or "").strip()) - - -def namespace_root() -> NavAddress: - """Scope handle for a namespace root (not a navigable row).""" - return NavAddress(NavLevel.NAMESPACE, "") - - -def is_dispatch_only_level(level: Optional[NavLevel]) -> bool: - """Document (and namespace) nodes may be entered, not collected as leaves.""" - return level in (NavLevel.NAMESPACE, NavLevel.DOCUMENT) - - -def _coerce_level(level: Any) -> Optional[NavLevel]: - if isinstance(level, NavLevel): - return level - if isinstance(level, str): - try: - return NavLevel(level) - except ValueError: - return None - return None - - -def _coerce_doc_id(value: Any) -> str: - if isinstance(value, str): - return value.strip() - return "" - - -def address_level(ts: Any, node_id: str) -> Optional[NavLevel]: - """Read level from the toolspace/provider registry when available.""" - sid = str(node_id or "").strip() - if not sid: - return NavLevel.NAMESPACE - fn = getattr(ts, "address_level", None) - if callable(fn): - level = _coerce_level(fn(sid)) - if level is not None: - return level - provider = getattr(ts, "_provider", None) - fn = getattr(provider, "address_level", None) - if callable(fn): - return _coerce_level(fn(sid)) - return None - - -def is_dispatch_only_node(ts: Any, node_id: str) -> bool: - """True when the node is document/namespace scoped (DISPATCH only).""" - return is_dispatch_only_level(address_level(ts, node_id)) - - -def owner_document(ts: Any, node_id: str, fallback: str = "") -> str: - """Owning ``document_id`` for a node, without parsing id strings.""" - sid = str(node_id or "").strip() - fb = str(fallback or "").strip() - if not sid: - return fb - - fn = getattr(ts, "owner_document", None) - if callable(fn): - got = _coerce_doc_id(fn(sid)) - if got: - return got - - provider = getattr(ts, "_provider", None) - fn = getattr(provider, "owner_document", None) - if callable(fn): - got = _coerce_doc_id(fn(sid)) - if got: - return got - - synth = getattr(ts, "_synthetic_doc_id", None) - if callable(synth): - got = _coerce_doc_id(synth(sid)) - if got: - return got - - idx = getattr(ts, "_idx", None) - if idx is not None: - loc = getattr(idx, "_node_to_doc_line", {}).get(sid) - if loc: - return str(loc[0]) - - return fb - - -def uses_document_nodes(ts: Any) -> bool: - """True when document ids themselves are map nodes (namespace provider). - - Checks the provider class (not a duck-typed instance attribute) so adapters - and test mocks do not look like namespace mode by accident. - """ - provider = getattr(ts, "_provider", None) - target = provider if provider is not None else ts - return callable(getattr(type(target), "document_ids", None)) - - -def next_dispatch_depth( - ts: Any, - *, - parent_doc_id: str = "", - parent_scope: str = "", - child_id: str, - depth: int, -) -> int: - """Child harvest depth after DISPATCH into ``child_id``. - - Namespace → document is depth-neutral (child starts at 0). Namespace → - section under a document starts at 1. Once already inside a document or - section scope, always ``depth + 1``. - """ - parent_sid = str(parent_scope or "").strip() - if parent_sid: - parent_level = address_level(ts, parent_sid) - if parent_level in (NavLevel.DOCUMENT, NavLevel.SECTION, NavLevel.CHUNK): - return int(depth) + 1 - - namespace_parent = uses_document_nodes(ts) and not str(parent_doc_id or "").strip() - if not namespace_parent: - return int(depth) + 1 - level = address_level(ts, child_id) - if level == NavLevel.DOCUMENT: - return 0 - if owner_document(ts, child_id, ""): - return 1 - return int(depth) + 1 diff --git a/deprecated/mapnav/nav/nav_agent.py b/deprecated/mapnav/nav/nav_agent.py deleted file mode 100644 index c7334ff9b..000000000 --- a/deprecated/mapnav/nav/nav_agent.py +++ /dev/null @@ -1,510 +0,0 @@ -from __future__ import annotations - -import logging -import os -import time -from typing import Any, List, Optional, Sequence, Tuple - -from ._compat import AgentStep, EpisodeResult -from ._compat import compose_answer_llm -from ._compat import HierarchicalTools -from ._compat import Chunk -from ._compat import ( - ToolSpace, -) -from .nav_address import ( - is_dispatch_only_node, - owner_document, - uses_document_nodes, -) -from .nav_compose import ( - evidence_owner_section_id, - pack_nav_evidence, - unit_score_for_evidence_chunk, -) -from .nav_types import ( - LegalAction, - NavConfig, - NavState, -) - -# Back-compat aliases for tests / callers. -_evidence_owner_section_id = evidence_owner_section_id -_unit_score_for_evidence_chunk = unit_score_for_evidence_chunk -_logger = logging.getLogger(__name__) - -# Large-scope title-only threshold = episode evidence budget x this mult. -_SCOPE_SUMMARY_BUDGET_MULT = 3.0 - - -def _chunks_to_retrieved_nodes(chunks: List[Chunk]) -> List[str]: - """Stable unit ids: prefer chunk/node id (Knowhere ``chunk_id``).""" - seen: set[str] = set() - out: List[str] = [] - for c in chunks: - node = str(getattr(c, "node_id", "") or "").strip() - if not node: - node = str(getattr(c, "section_id", "") or "").strip() - if not node or node in seen: - continue - seen.add(node) - out.append(node) - return out - - -def _dedupe_scored(scored: List[Tuple[Chunk, float]]) -> List[Tuple[Chunk, float]]: - from .nav_compose import dedupe_scored - - return dedupe_scored(scored) - - -def _add_scored(state: NavState, scored: List[Tuple[Chunk, float]]) -> int: - added = 0 - for c, score in scored: - if c.node_id in state.collected_ids: - continue - state.collected_ids.add(c.node_id) - state.collected.append((c, float(score))) - added += 1 - return added - - -def _resolve_action_doc_id( - action_or_sid: Any, - state: NavState, - ts: Any = None, -) -> str: - """Owning document_id for hydrate; never a corpus/namespace sentinel.""" - if hasattr(action_or_sid, "section_id"): - sid = str(getattr(action_or_sid, "section_id", "") or "").strip() - else: - sid = str(action_or_sid or "").strip() - resolved = owner_document(ts, sid, "") if ts is not None else "" - if resolved: - return resolved - if state.doc_id: - return state.doc_id - return "" - - -def _purge_descendant_evidence( - ts: ToolSpace, - state: NavState, - parent_sid: str, -) -> int: - """Drop standalone evidence owned by proper descendants of parent_sid. - - When COLLECT parent after COLLECT child (e.g. L93 then L92), child chunks - are absorbed into the parent hydrate and must not keep a separate bag slot. - """ - sid = str(parent_sid or "").strip() - if not sid: - return 0 - doc = _resolve_action_doc_id(sid, state, ts) - relations = getattr(ts, "section_relation_ids", None) - if callable(relations): - _anc, descendants = relations(sid, doc) - else: - descendants = _section_and_descendants(ts, sid, doc) - descendants = {str(x).strip() for x in (descendants or set()) if str(x).strip()} - descendants.discard(sid) - if not descendants: - return 0 - - kept: List[Tuple[Chunk, float]] = [] - removed = 0 - for chunk, score in list(state.collected): - owner = _evidence_owner_section_id(chunk) - if owner in descendants: - state.collected_ids.discard(chunk.node_id) - removed += 1 - continue - kept.append((chunk, score)) - state.collected = kept - return removed - - -def _env_enabled(name: str, default: str = "1") -> bool: - return os.environ.get(name, default).strip().lower() not in {"0", "false", "no", "off"} - - -def _line_order(pool: List[Chunk]) -> List[Chunk]: - return sorted(pool, key=lambda c: (min(c.line_ids or (10**9,)), c.node_id)) - - -def _collect_subtree(ts: ToolSpace, action: LegalAction, state: NavState, config: NavConfig) -> List[Tuple[Chunk, float]]: - """Hydrate ``section_id ∪ descendants`` in document order. - - Final size is controlled by compose ``budget_chars`` progressive trim. - """ - sid = action.section_id - if not sid: - return [] - if is_dispatch_only_node(ts, sid): - return [] - doc = _resolve_action_doc_id(action, state, ts) - if not doc: - return [] - materialize = getattr(ts, "_materialize_leaf_path_chunks", None) - if not callable(materialize): - return [] - pool = list(materialize(sid, doc)) - if not pool: - return [] - return _collect_in_doc_order(pool, config) - - -def _mark_collected_branch( - ts: ToolSpace, action: LegalAction, state: NavState, added: int -) -> dict[str, Any]: - """On successful COLLECT: mark sid ∪ descendants as collected (removed from map). - - Replaces the old covered/collected split — one set only. - """ - sid = str(action.section_id or "").strip() - if not sid or not _env_enabled("NAV_FILTER_COLLECTED_SECTIONS"): - return {} - doc = _resolve_action_doc_id(action, state, ts) - materialize = getattr(ts, "_materialize_leaf_path_chunks", None) - relations = getattr(ts, "section_relation_ids", None) - pool = list(materialize(sid, doc)) if callable(materialize) and doc else [] - if callable(relations) and doc: - ancestors, descendants = relations(sid, doc) - else: - ancestors, descendants = set(), _section_and_descendants(ts, sid, doc) - descendants = {str(x).strip() for x in (descendants or set()) if str(x).strip()} - descendants.add(sid) - - is_full = bool(pool) and all(chunk.node_id in state.collected_ids for chunk in pool) - if added > 0: - state.collected_section_ids.update(descendants) - if len(pool) > 1: - state.blocked_collect_section_ids.update(ancestors) - return { - "collect_full": is_full, - "branch_selected": added > 0, - "n_collected_sections": len(state.collected_section_ids), - "n_blocked_ancestor_collects": len(state.blocked_collect_section_ids), - } - - -# Back-compat name used by older tests; prefer _mark_collected_branch. -_update_collect_coverage = _mark_collected_branch - - -def _collect_in_doc_order( - pool: List[Chunk], - config: NavConfig, -) -> List[Tuple[Chunk, float]]: - """Hydrate the complete branch in document order; COMPOSE truncates later.""" - if not pool: - return [] - base = float(config.read_score_bonus) - return [(chunk, base) for chunk in _line_order(pool)] - - -def _direct_child_ids(ts: ToolSpace, section_id: str, doc_id: str) -> List[str]: - sid = (section_id or "").strip() - if not sid: - return [] - children_fn = getattr(ts, "_children_for_section_path", None) - rows: List[Any] = [] - if callable(children_fn): - try: - rows = list(children_fn(sid, doc_id) or []) - except Exception: - rows = [] - if not rows: - try: - st = ts.get_structure(sid) - rows = list(st.get("children") or []) - except Exception: - return [] - out: List[str] = [] - seen: set[str] = set() - for row in rows: - if not isinstance(row, dict): - continue - child_id = str(row.get("section_id") or "").strip() - if child_id and child_id not in seen: - seen.add(child_id) - out.append(child_id) - return out - - -def _section_and_descendants(ts: ToolSpace, section_id: str, doc_id: str) -> set[str]: - """Return {section_id} ∪ descendants (best-effort).""" - sid = (section_id or "").strip() - if not sid: - return set() - relations = getattr(ts, "section_relation_ids", None) - if callable(relations): - try: - _anc, desc = relations(sid, doc_id) - out = {str(x).strip() for x in (desc or set()) if str(x).strip()} - out.add(sid) - return out - except Exception: - pass - out: set[str] = {sid} - stack = [sid] - while stack: - cur = stack.pop() - for child in _direct_child_ids(ts, cur, doc_id): - if child not in out: - out.add(child) - stack.append(child) - return out - - -def _episode_stop_reason(steps: Sequence[Any]) -> str: - """Code-derived episode stop label from budget state / step meta (not LLM).""" - from .nav_token_budget import nav_token_budget_exhausted - - if nav_token_budget_exhausted(): - return "token_limit" - for step in steps: - detail = getattr(step, "detail", None) or {} - if not isinstance(detail, dict): - continue - sr = str(detail.get("stop_reason") or "").strip() - if sr: - return sr - if str(detail.get("reason") or "").strip() == "token_limit": - return "token_limit" - return "completed" - - -def run_nav_episode( - tools: Optional[HierarchicalTools], - query: str, - *, - doc_id: Optional[str] = None, - corpus_doc_ids: Optional[Sequence[str]] = None, - budget_chars: int, - task_type: str = "unknown", - compose_format_constraints: str = "", - compose_answer: bool = True, - policy: str = "rule", - config: Optional[NavConfig] = None, - toolspace: Optional[Any] = None, -) -> EpisodeResult: - from .nav_llm import nav_llm_runtime - from .nav_token_budget import nav_token_episode - - cfg = config or NavConfig(policy="llm") - with nav_token_episode(token_limit=int(getattr(cfg, "token_limit", 0) or 0)): - with nav_llm_runtime( - planner_thinking=str(getattr(cfg, "planner_thinking", "") or ""), - planner_think_max_tokens=int( - getattr(cfg, "planner_think_max_tokens", 0) or 0 - ), - ): - return _run_nav_episode_body( - tools, - query, - doc_id=doc_id, - corpus_doc_ids=corpus_doc_ids, - budget_chars=budget_chars, - task_type=task_type, - compose_format_constraints=compose_format_constraints, - compose_answer=compose_answer, - policy=policy, - config=cfg, - toolspace=toolspace, - ) - - -def _run_nav_episode_body( - tools: Optional[HierarchicalTools], - query: str, - *, - doc_id: Optional[str] = None, - corpus_doc_ids: Optional[Sequence[str]] = None, - budget_chars: int, - task_type: str = "unknown", - compose_format_constraints: str = "", - compose_answer: bool = True, - policy: str = "rule", - config: Optional[NavConfig] = None, - toolspace: Optional[Any] = None, -) -> EpisodeResult: - from .nav_token_budget import stamp_step_detail - - corpus_ids = [ - str(d).strip() - for d in (corpus_doc_ids or []) - if str(d).strip() - ] - episode_doc = str(doc_id or "").strip() - if not episode_doc and not corpus_ids: - raise ValueError( - "Nav Agent requires doc_id or non-empty corpus_doc_ids " - "(eval entry points always pass corpus_doc_ids for task_corpus)" - ) - from ._compat import load_llm_env, require_llm_env # type: ignore - - load_llm_env() - require_llm_env(context="Nav Agent") - cfg = config or NavConfig(policy="llm") - if cfg.llm_max_tokens < 256: - cfg.llm_max_tokens = 256 - nav_policy = (policy or cfg.policy or "llm").strip().lower() - if nav_policy != "llm": - raise ValueError( - f"Nav Agent 仅支持 llm 策略（须配置 DS_KEY 或 OPENAI_API_KEY）；" - f"收到 policy={policy!r}。" - "请设置 --nav-policy llm 或删除 NAV_POLICY=rule。" - ) - cfg.policy = "llm" - # Tie the large-scope title-only threshold to the real evidence budget - # (budget_chars x mult): a scope whose full summary map would dwarf the final - # evidence budget is shown title-only, nudging DISPATCH over broad COLLECT. - if int(budget_chars) > 0: - cfg.scope_inline_summary_char_limit = max( - 1, int(budget_chars * _SCOPE_SUMMARY_BUDGET_MULT) - ) - retrieval_t0 = time.perf_counter() - if toolspace is not None: - ts = toolspace - elif tools is not None: - ts = ToolSpace(tools, corpus_doc_ids=corpus_ids or None) - else: - raise ValueError("Nav Agent requires either tools or an injected toolspace") - - # Namespace / multi-doc: document ids are map nodes; empty scope is the root. - namespace_mode = bool(corpus_ids) and uses_document_nodes(ts) - if namespace_mode: - if not corpus_ids: - corpus_ids = list(ts.document_ids()) - episode_doc = "" - - state = NavState(doc_id=episode_doc, query=query, task_type=task_type) - steps: List[AgentStep] = [] - if namespace_mode: - section_ids = list(ts.sections_for_doc("")) - else: - section_ids = ts.sections_for_doc(episode_doc) - - from .nav_plan import plan_query - - # Planner is query-only: no pre-lit map. Score/light after the plan exists. - plan_t0 = time.perf_counter() - retrieval_plan = plan_query(state, cfg) - state.retrieval_plan = retrieval_plan - steps.append( - AgentStep( - step_idx=len(steps) + 1, - action="query_plan", - detail=stamp_step_detail({ - "fallback": bool(retrieval_plan.fallback), - "n_subgoals": len(retrieval_plan.subgoals), - "reason": retrieval_plan.reason, - "plan": retrieval_plan.to_dict(), - "seconds": time.perf_counter() - plan_t0, - }, t0=plan_t0), - ) - ) - _logger.info( - "retrieval mapnav phase=planner seconds=%.3f subgoals=%d", - time.perf_counter() - plan_t0, - len(retrieval_plan.subgoals), - ) - - # Checklist: wave orchestration. Every episode runs plan + harvest + control. - if state.retrieval_plan is None: - from .nav_plan import fallback_plan - - state.retrieval_plan = fallback_plan(query, reason="missing_plan") - from .nav_orchestrate import execute_plan, seed_episode_map_scores_from_plan - - seed_episode_map_scores_from_plan(ts, state, cfg, state.retrieval_plan) - - orch_t0 = time.perf_counter() - orch_detail = execute_plan( - ts, state, cfg, steps_out=steps, episode_query=query - ) - orch_detail = dict(orch_detail or {}) - orch_detail["seconds"] = time.perf_counter() - orch_t0 - steps.append( - AgentStep( - step_idx=len(steps) + 1, - action="plan_orchestrate", - detail=stamp_step_detail(orch_detail, t0=orch_t0), - ) - ) - _logger.info( - "retrieval mapnav phase=orchestration seconds=%.3f waves=%d", - time.perf_counter() - orch_t0, - len(orch_detail.get("waves", [])), - ) - - evidence_started = time.perf_counter() - fill = pack_nav_evidence( - _dedupe_scored(list(state.collected)), - ts, - state, - cfg, - budget_chars=budget_chars, - ) - _logger.info( - "retrieval mapnav phase=evidence_pack seconds=%.3f chunks=%d", - time.perf_counter() - evidence_started, - len(fill.kept_chunks), - ) - scored_chunks = list(fill.scored_chunks) - retrieval_seconds = time.perf_counter() - retrieval_t0 - composed = "" - compose_seconds = 0.0 - if compose_answer: - compose_t0 = time.perf_counter() - max_ans = min(1024, max(256, int(budget_chars))) - extra_mh_constraint = "" - if (task_type or "").strip().lower() == "multi_hop": - extra_mh_constraint = ( - "multi_hop 约束：fact_1 与 fact_2 必须分别覆盖两跳信息，" - "final_answer 必须整合二者，任一缺失视为不完整。" - ) - fc = compose_format_constraints - if extra_mh_constraint: - fc = (f"{fc}\n{extra_mh_constraint}" if fc else extra_mh_constraint) - composed = compose_answer_llm( - query, - task_type=task_type or "niche_fact", - evidence_text=fill.evidence_text or "", - max_answer_chars=max_ans, - budget_chars=int(budget_chars), - format_constraints=fc, - ) - compose_seconds = time.perf_counter() - compose_t0 - steps.append( - AgentStep( - step_idx=len(steps) + 1, - action="compose_answer", - detail=stamp_step_detail({ - "evidence_chars": fill.evidence_chars_actual, - "n_chunks_kept": fill.n_chunks_kept, - "truncated_last": fill.truncated_last, - }, t0=compose_t0), - ) - ) - - return EpisodeResult( - representation=f"hierarchical_nav_{cfg.policy}", - steps=steps, - scored_chunks=scored_chunks, - kept_chunks=fill.kept_chunks, - evidence_text=fill.evidence_text, - evidence_chars_actual=fill.evidence_chars_actual, - retrieved_nodes=_chunks_to_retrieved_nodes(list(fill.kept_chunks)), - composed_answer=composed, - section_ids=list(section_ids), - trajectory_length=len(steps), - truncated_last=fill.truncated_last, - phase_timings={ - "retrieval_framework_seconds": retrieval_seconds, - "compose_seconds": compose_seconds, - "online_response_seconds": retrieval_seconds + compose_seconds, - }, - stop_reason=_episode_stop_reason(steps), - ) diff --git a/deprecated/mapnav/nav/nav_assets.py b/deprecated/mapnav/nav/nav_assets.py deleted file mode 100644 index 60fccba7f..000000000 --- a/deprecated/mapnav/nav/nav_assets.py +++ /dev/null @@ -1,758 +0,0 @@ -"""Asset SEARCH channel aligned with Knowhere ``assets.py``. - -Pipeline (same shape as production ``search_assets_step``): - -1. Document-bound scope gather by ``chunk_type`` (``asset_filter_step``). -2. Asset inspector: tables → text LLM; images → VLM when available, else - text LLM fallback (and VLM failure also falls back). -3. Matched assets only are added to evidence; inspector context is injected - into the next harvest observation (``NavState.asset_observation_context``). - -Namespace / unbound scope is skipped — never scan the whole corpus. -""" - -from __future__ import annotations - -import json -import re -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple - -from .nav_address import NavLevel, address_level, owner_document -from .nav_knowhere import is_root_section - -# Knowhere DocumentChunk.chunk_type values + harvest JSON kind aliases. -_KIND_TO_CHUNK_TYPE = { - "image": "image", - "images": "image", - "table": "table", - "tables": "table", -} - -# Soft prompt char budget for the text inspector (episode budget is separate). -_ASSET_FILTER_PROMPT_CHAR_BUDGET = 24000 - -NavVlmBackend = Callable[..., Any] -_vlm_backend: Optional[NavVlmBackend] = None - - -def set_nav_vlm_backend(fn: Optional[NavVlmBackend]) -> None: - """Install or clear an injected VLM backend (tests / Knowhere vlm_fn).""" - global _vlm_backend - _vlm_backend = fn - - -def get_nav_vlm_backend() -> Optional[NavVlmBackend]: - return _vlm_backend - - -def asset_chunk_type(kind: str) -> Optional[str]: - """Map harvest ``kind`` to Knowhere ``chunk_type``, or None if unknown.""" - return _KIND_TO_CHUNK_TYPE.get(str(kind or "").strip().lower()) - - -def parse_search_assets(raw: Any) -> List[Dict[str, str]]: - """Normalize harvest ``search_assets`` entries to ``{kind,query,scope,asset_type}``.""" - if not isinstance(raw, list): - return [] - out: List[Dict[str, str]] = [] - for item in raw: - if not isinstance(item, dict): - continue - kind = str(item.get("kind") or "").strip() - asset_type = asset_chunk_type(kind) - if not asset_type: - continue - query = str(item.get("query") or "").strip() - scope = str(item.get("scope") or "").strip() - out.append( - { - "kind": kind, - "query": query, - "scope": scope, - "asset_type": asset_type, - } - ) - return out - - -def _provider(ts: Any) -> Any: - return getattr(ts, "_provider", None) - - -def resolve_asset_search_scope( - ts: Any, - *, - requested_scope: Optional[str], - default_scope: Optional[str], - fallback_doc_id: str = "", -) -> Tuple[str, str, Optional[str]]: - """Bind asset SEARCH to one document. - - Returns ``(effective_scope, doc_id, skip_reason)``. - """ - raw = str(requested_scope or "").strip() or str(default_scope or "").strip() - if not raw: - return "", "", "skipped_no_document_scope" - - level = address_level(ts, raw) - if level == NavLevel.NAMESPACE: - return "", "", "skipped_namespace_scope" - - doc = owner_document(ts, raw, fallback_doc_id) - if not doc and level == NavLevel.DOCUMENT: - doc = raw - doc = str(doc or "").strip() - if not doc: - return "", "", "skipped_no_document_scope" - - if level in (NavLevel.SECTION, NavLevel.CHUNK): - owner = owner_document(ts, raw, doc) - if owner and owner != doc: - return "", "", "skipped_cross_document_scope" - - return raw, doc, None - - -def _section_ids_in_document(ts: Any, doc_id: str) -> List[str]: - doc = str(doc_id or "").strip() - if not doc: - return [] - - relations = getattr(ts, "section_relation_ids", None) - if callable(relations): - _anc, descendants = relations(doc, doc) - got = [str(x) for x in (descendants or ()) if str(x).strip()] - if got: - return list(dict.fromkeys(got)) - - provider = _provider(ts) - if provider is not None and callable(getattr(provider, "all_section_ids", None)): - out = [ - sid - for sid in (str(x) for x in provider.all_section_ids() if str(x).strip()) - if owner_document(ts, sid, doc) == doc - ] - if out: - return list(dict.fromkeys(out)) - - roots_fn = getattr(ts, "sections_for_doc", None) - if callable(roots_fn): - out: List[str] = [] - for root in (str(x) for x in (roots_fn(doc) or ()) if str(x).strip()): - if owner_document(ts, root, doc) != doc: - continue - out.extend(_section_ids_under_bound_scope(ts, root, doc)) - return list(dict.fromkeys(out)) - return [] - - -def _section_ids_under_bound_scope(ts: Any, scope: str, doc_id: str) -> List[str]: - sid = str(scope or "").strip() - doc = str(doc_id or "").strip() - if not sid or not doc: - return [] - - level = address_level(ts, sid) - if level == NavLevel.DOCUMENT or sid == doc: - return _section_ids_in_document(ts, doc) - - provider = _provider(ts) - out: List[str] = [sid] - relations = getattr(ts, "section_relation_ids", None) - if callable(relations): - _anc, descendants = relations(sid, doc) - out.extend(str(x) for x in (descendants or ()) if str(x).strip()) - elif provider is not None and callable(getattr(provider, "relations", None)): - _anc, descendants = provider.relations(sid) - out.extend(str(x) for x in (descendants or ()) if str(x).strip()) - - return [ - s for s in dict.fromkeys(out) if owner_document(ts, s, doc) == doc or s == sid - ] - - -def gather_scoped_asset_candidates( - ts: Any, - *, - asset_type: str, - scope: Optional[str], - doc_id: str = "", -) -> List[Dict[str, Any]]: - """Units under a document-bound scope with matching ``chunk_type``.""" - from ._compat import Chunk # type: ignore - - wanted = str(asset_type or "").strip().lower() - if wanted not in {"image", "table"}: - return [] - - resolved_doc = str(doc_id or "").strip() - sid = str(scope or "").strip() - if not sid or not resolved_doc: - return [] - - provider = _provider(ts) - self_units = getattr(provider, "self_units", None) if provider is not None else None - unit_text = getattr(provider, "unit_text", None) if provider is not None else None - if not callable(self_units) or not callable(unit_text): - return [] - - section_ids = _section_ids_under_bound_scope(ts, sid, resolved_doc) - out: List[Dict[str, Any]] = [] - seen: Set[str] = set() - section_path_fn = getattr(provider, "section_path", None) - for section_id in section_ids: - owner = owner_document(ts, section_id, resolved_doc) - if owner and owner != resolved_doc: - continue - # Align with Knowhere: never treat Root as an asset owner. - if is_root_section(provider, section_id): - continue - owner_path = ( - str(section_path_fn(section_id) or "").strip() - if callable(section_path_fn) - else str(section_id) - ) - for unit in self_units(section_id) or (): - if str(getattr(unit, "chunk_type", "") or "").strip().lower() != wanted: - continue - chunk_id = str(getattr(unit, "chunk_id", "") or "").strip() - if not chunk_id or chunk_id in seen: - continue - text = str(unit_text(unit) or "").strip() - if not text: - continue - seen.add(chunk_id) - meta = dict(getattr(unit, "metadata", None) or {}) - file_path = str( - getattr(unit, "file_path", "") - or meta.get("file_path") - or getattr(unit, "source_chunk_path", "") - or getattr(unit, "content", "") - or "" - ).strip() - summary = str(meta.get("summary") or "").strip() - chunk = Chunk( - node_id=chunk_id, - doc_id=owner or resolved_doc, - text=text, - line_ids=(int(getattr(unit, "sort_order", 0) or 0),), - section_id=section_id, - ) - out.append( - { - "chunk_id": chunk_id, - "chunk_type": wanted, - "content": str(getattr(unit, "content", "") or ""), - "file_path": file_path, - "section_id": section_id, - "section_path": owner_path or str(section_id), - "owner_section_path": owner_path or str(section_id), - "summary": summary, - "chunk_metadata": meta, - "display_text": text, - "chunk": chunk, - "url": str(meta.get("url") or meta.get("presigned_url") or "").strip(), - } - ) - out.sort(key=lambda a: (min(a["chunk"].line_ids or (0,)), a["chunk_id"])) - return out - - -def gather_scoped_asset_chunks( - ts: Any, - *, - asset_type: str, - scope: Optional[str], - doc_id: str = "", -) -> List[Any]: - """Back-compat: Chunk list for scope gather (no inspector).""" - return [ - c["chunk"] - for c in gather_scoped_asset_candidates( - ts, asset_type=asset_type, scope=scope, doc_id=doc_id - ) - ] - - - -def _markdown_cell(value: str) -> str: - return ( - str(value or "") - .replace("\n", " ") - .replace("\r", " ") - .replace("|", "\\|") - .strip() - ) - - -def _format_asset_candidates_table(candidates: List[Dict[str, str]]) -> str: - lines = [ - "| ID | File | Summary / content signal |", - "|---|---|---|", - ] - for candidate in candidates: - lines.append( - "| " - + " | ".join( - [ - _markdown_cell(candidate.get("id", "")), - _markdown_cell(candidate.get("file", "")), - _markdown_cell(candidate.get("desc", "")), - ] - ) - + " |" - ) - return "\n".join(lines) - - -def _format_asset_filter_prompt( - query: str, - asset_type: str, - candidates: List[Dict[str, str]], -) -> str: - type_label = "images" if asset_type == "image" else "tables" - items_text = _format_asset_candidates_table(candidates) - example_id = candidates[0]["id"] if candidates else ("I1" if asset_type == "image" else "T1") - return ( - f"You are an asset relevance filter.\n\n" - f"Original user query: {query}\n\n" - f"Below are {len(candidates)} {type_label} from a document. " - f"Select ONLY assets that directly satisfy the user's query.\n\n" - f"Selection policy:\n" - f"- Match the requested asset type and the requested subject. " - f"Being an image/chart/table is not enough.\n" - f"- Do not select assets only because they belong to the same broad " - f"domain as the query.\n" - f"- Do not broaden specific market, instrument, company, metric, or " - f"entity terms. Neighboring topics are not matches unless the candidate " - f"explicitly connects them to the requested subject.\n" - f"- Treat words like \"all\" as all relevant assets, not all visible " - f"candidates.\n" - f"- If the file name, summary, or content signal does not " - f"directly support relevance, leave it out.\n" - f"- If uncertain, do not select the asset.\n\n" - f"=== Candidate {type_label.title()} ===\n{items_text}\n=== End ===\n\n" - f"Return ONLY a JSON array of matching row IDs, e.g.: " - f'["{example_id}"]\n' - f"If none are relevant, return an empty array: []\n" - f"Do not include any explanation." - ) - - -def _fit_text_to_char_budget(text: str, char_budget: int) -> str: - text = text.strip() - if not text or len(text) <= char_budget: - return text - return text[: max(0, char_budget)].rstrip() + "…" - - -def _project_assets_for_text_filter( - *, - asset_type: str, - assets: Sequence[Dict[str, Any]], -) -> Tuple[List[Dict[str, str]], Set[str], Dict[str, str]]: - projected: List[Dict[str, str]] = [] - valid_ids: Set[str] = set() - id_to_chunk_id: Dict[str, str] = {} - for index, asset in enumerate(assets, start=1): - chunk_id = str(asset.get("chunk_id") or "") - if not chunk_id: - continue - row_id = f"I{index}" if asset_type == "image" else f"T{index}" - summary = str(asset.get("summary") or "").strip() - file_path = str(asset.get("file_path") or "") - content = str(asset.get("content") or "").strip() - display = str(asset.get("display_text") or "").strip() - description = summary or display or (content if asset_type == "table" else "") - projected.append({"id": row_id, "file": file_path, "desc": description}) - valid_ids.add(row_id) - id_to_chunk_id[row_id] = chunk_id - - if not projected: - return projected, valid_ids, id_to_chunk_id - - prompt = _format_asset_filter_prompt("q", asset_type, projected) - if len(prompt) <= _ASSET_FILTER_PROMPT_CHAR_BUDGET: - return projected, valid_ids, id_to_chunk_id - - # Compact descriptions so the full prompt fits the soft envelope. - overhead = len(_format_asset_filter_prompt("q", asset_type, [ - {"id": item["id"], "file": item["file"], "desc": ""} for item in projected - ])) - desc_budget = max(_ASSET_FILTER_PROMPT_CHAR_BUDGET - overhead, len(projected)) - per_item = max(desc_budget // len(projected), 8) - compacted = [ - { - "id": item["id"], - "file": item["file"], - "desc": _fit_text_to_char_budget(item["desc"], per_item), - } - for item in projected - ] - return compacted, valid_ids, id_to_chunk_id - - -def _parse_asset_filter_response(text: str, valid_ids: Set[str]) -> List[str]: - text = (text or "").strip() - try: - result = json.loads(text) - if isinstance(result, list): - return [str(item) for item in result if str(item) in valid_ids] - except (ValueError, json.JSONDecodeError): - pass - - fence_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL) - if fence_match: - try: - result = json.loads(fence_match.group(1).strip()) - if isinstance(result, list): - return [str(item) for item in result if str(item) in valid_ids] - except (ValueError, json.JSONDecodeError): - pass - - bracket_match = re.search(r"\[.*?\]", text, re.DOTALL) - if bracket_match: - try: - result = json.loads(bracket_match.group()) - if isinstance(result, list): - return [str(item) for item in result if str(item) in valid_ids] - except (ValueError, json.JSONDecodeError): - pass - return [] - - -def _search_assets_via_text_llm( - *, - query: str, - asset_type: str, - assets: Sequence[Dict[str, Any]], - config: Any, -) -> List[str]: - from .nav_llm import nav_chat, resolve_nav_model - - projected, valid_ids, id_to_chunk_id = _project_assets_for_text_filter( - asset_type=asset_type, assets=assets - ) - if not projected: - return [] - - prompt = _format_asset_filter_prompt(query, asset_type, projected) - model = resolve_nav_model( - model=str(getattr(config, "llm_model", "") or ""), - model_env="NAV_LLM_MODEL", - fallback_envs=(), - ) - try: - cached = nav_chat( - purpose="nav_asset_filter_v1", - model=model, - messages=[{"role": "user", "content": prompt}], - temperature=0.0, - max_tokens=max(256, int(getattr(config, "llm_max_tokens", 256) or 256)), - context="Nav Asset Filter", - usage_tag="nav_asset_filter", - ) - text = str(cached.get("content") or "") - # Some models wrap the array in {"ids":[...]}; accept either. - selected = _parse_asset_filter_response(text, valid_ids) - if not selected: - try: - obj = json.loads(text) - if isinstance(obj, dict): - for key in ("ids", "selected", "matches", "result"): - if isinstance(obj.get(key), list): - selected = [ - str(x) for x in obj[key] if str(x) in valid_ids - ] - break - except (ValueError, json.JSONDecodeError): - pass - return [ - id_to_chunk_id[row_id] - for row_id in selected - if row_id in id_to_chunk_id - ] - except Exception: - return [] - - -def _resolve_image_url(asset: Dict[str, Any]) -> str: - url = str(asset.get("url") or "").strip() - if url: - return url - meta = asset.get("chunk_metadata") or {} - if isinstance(meta, dict): - for key in ("url", "presigned_url", "image_url"): - got = str(meta.get(key) or "").strip() - if got: - return got - return "" - - -def _search_images_via_vlm( - *, - query: str, - assets: Sequence[Dict[str, Any]], - vlm_fn: NavVlmBackend, -) -> Tuple[List[str], Optional[str]]: - """VLM image filter when URLs exist; otherwise signal fallback.""" - candidates: List[Tuple[str, str, str]] = [] - valid_ids: Set[str] = set() - id_to_chunk_id: Dict[str, str] = {} - for index, asset in enumerate(assets, start=1): - chunk_id = str(asset.get("chunk_id") or "") - url = _resolve_image_url(asset) - if not chunk_id or not url: - continue - row_id = f"I{index}" - candidates.append((row_id, str(asset.get("file_path") or ""), url)) - valid_ids.add(row_id) - id_to_chunk_id[row_id] = chunk_id - - if not candidates: - return [], "no_image_urls" - - content_parts: List[Dict[str, Any]] = [ - { - "type": "text", - "text": ( - f"You are an image relevance filter.\n\n" - f"Original user query: {query}\n\n" - f"Below are {len(candidates)} images from a document. " - f"Look at each image and select ONLY images that directly " - f"satisfy the user's query.\n\n" - f"Return ONLY a JSON array of matching row IDs. " - f"If none are relevant, return [].\n" - ), - } - ] - for row_id, file_path, url in candidates: - content_parts.append({"type": "text", "text": f"\n{row_id} | {file_path}\n"}) - content_parts.append({"type": "image_url", "image_url": {"url": url}}) - - try: - response = vlm_fn( - purpose="nav_asset_vlm_v1", - messages=[{"role": "user", "content": content_parts}], - ) - if isinstance(response, dict): - text = str(response.get("content") or "") - else: - text = str(response or "") - selected = _parse_asset_filter_response(text, valid_ids) - return [ - id_to_chunk_id[row_id] - for row_id in selected - if row_id in id_to_chunk_id - ], None - except Exception as exc: - return [], f"vlm_failed:{type(exc).__name__}" - - -def search_assets_step( - *, - query: str, - asset_type: str, - candidates: Sequence[Dict[str, Any]], - config: Any, -) -> Dict[str, Any]: - """Knowhere-aligned inspector over already-gathered candidates.""" - assets = list(candidates) - if not assets: - return { - "status": "empty", - "status_detail": "", - "matched_assets": [], - "candidate_count": 0, - } - - asset_by_id = { - str(a.get("chunk_id") or ""): a - for a in assets - if str(a.get("chunk_id") or "") - } - status_detail = "" - selected_ids: List[str] = [] - - if asset_type == "image": - vlm_fn = get_nav_vlm_backend() - if vlm_fn is None: - selected_ids = _search_assets_via_text_llm( - query=query, asset_type=asset_type, assets=assets, config=config - ) - status = "fallback_matched" if selected_ids else "fallback_empty" - status_detail = "vlm_unavailable_text_fallback" - else: - selected_ids, vlm_error = _search_images_via_vlm( - query=query, assets=assets, vlm_fn=vlm_fn - ) - if vlm_error: - selected_ids = _search_assets_via_text_llm( - query=query, asset_type=asset_type, assets=assets, config=config - ) - status = "fallback_matched" if selected_ids else "fallback_empty" - status_detail = f"vlm_failed_text_fallback:{vlm_error}" - else: - status = "matched" if selected_ids else "empty" - else: - selected_ids = _search_assets_via_text_llm( - query=query, asset_type=asset_type, assets=assets, config=config - ) - status = "matched" if selected_ids else "empty" - - matched = [asset_by_id[cid] for cid in selected_ids if cid in asset_by_id] - return { - "status": status, - "status_detail": status_detail, - "matched_assets": matched, - "candidate_count": len(asset_by_id), - } - - -def format_asset_observation_context( - *, - kind: str, - asset_type: str, - search_query: str, - matched_assets: Sequence[Dict[str, Any]], - status: str = "empty", - status_detail: str = "", -) -> str: - """Next-step observation block (Knowhere ``_format_asset_context``).""" - tool_name = "SEARCH_IMAGES" if asset_type == "image" else "SEARCH_TABLES" - kind_norm = str(kind or "").strip().lower() - if kind_norm in {"search_images", "images", "image"}: - tool_name = "SEARCH_IMAGES" - elif kind_norm in {"search_tables", "tables", "table"}: - tool_name = "SEARCH_TABLES" - if not matched_assets: - detail = f" Status detail: {status_detail}." if status_detail else "" - return ( - f"=== {tool_name} Results ===\n" - f"No matching {asset_type}s found for \"{search_query}\" " - f"(status={status}).{detail}\n" - f"=== End {tool_name} Results ===" - ) - lines = [ - f"=== {tool_name} Results ===", - f'Found {len(matched_assets)} matching {asset_type}s for "{search_query}".', - "Matched assets are available as asset evidence.", - ] - for i, asset in enumerate(matched_assets, 1): - file_path = str(asset.get("file_path") or asset.get("chunk_id") or "") - lines.append(f" {i}. {file_path}") - owners = [] - seen: Set[str] = set() - for asset in matched_assets: - owner = str( - asset.get("owner_section_path") or asset.get("section_path") or "" - ).strip() - if owner and owner not in seen: - seen.add(owner) - owners.append(owner) - if owners: - lines.append("Owner sections with matching assets:") - for owner in owners: - lines.append(f' - "{owner}"') - lines.append( - "Use these asset results and owner sections to decide collect, " - "finish, or further navigation." - ) - lines.append(f"=== End {tool_name} Results ===") - return "\n".join(lines) - - -def apply_search_assets( - ts: Any, - state: Any, - config: Any, - *, - requests: Sequence[Dict[str, str]], - default_scope: Optional[str], -) -> Tuple[int, List[Dict[str, Any]]]: - """Inspect scoped assets; add matches to evidence; inject observation context.""" - from .nav_agent import _add_scored - - bonus = float(getattr(config, "read_score_bonus", 0.0) or 0.0) - trace: List[Dict[str, Any]] = [] - total = 0 - fallback_doc = str(getattr(state, "doc_id", "") or "") - context_blocks: List[str] = [] - - for req in requests: - asset_type = str(req.get("asset_type") or "").strip() - query = str(req.get("query") or "").strip() or str( - getattr(state, "query", "") or "" - ) - requested = str(req.get("scope") or "").strip() - scope, doc, skip_reason = resolve_asset_search_scope( - ts, - requested_scope=requested, - default_scope=default_scope, - fallback_doc_id=fallback_doc, - ) - if skip_reason: - trace.append( - { - "kind": req.get("kind", ""), - "asset_type": asset_type, - "query": query, - "scope": requested or (default_scope or ""), - "doc_id": "", - "n_candidates": 0, - "n_matched": 0, - "n_added": 0, - "status": "skipped", - "skip_reason": skip_reason, - } - ) - continue - - candidates = gather_scoped_asset_candidates( - ts, asset_type=asset_type, scope=scope, doc_id=doc - ) - result = search_assets_step( - query=query, - asset_type=asset_type, - candidates=candidates, - config=config, - ) - matched = list(result.get("matched_assets") or []) - scored = [(a["chunk"], bonus) for a in matched if a.get("chunk") is not None] - added = _add_scored(state, scored) if scored else 0 - total += added - - block = format_asset_observation_context( - kind=str(req.get("kind") or ""), - asset_type=asset_type, - search_query=query, - matched_assets=matched, - status=str(result.get("status") or "empty"), - status_detail=str(result.get("status_detail") or ""), - ) - context_blocks.append(block) - - trace.append( - { - "kind": req.get("kind", ""), - "asset_type": asset_type, - "query": query, - "scope": scope, - "doc_id": doc, - "n_candidates": int(result.get("candidate_count") or len(candidates)), - "n_matched": len(matched), - "n_added": added, - "status": result.get("status", ""), - "status_detail": result.get("status_detail", ""), - "matched_chunk_ids": [str(a.get("chunk_id") or "") for a in matched], - } - ) - - if context_blocks: - prev = str(getattr(state, "asset_observation_context", "") or "").strip() - joined = "\n\n".join(context_blocks) - state.asset_observation_context = ( - f"{prev}\n\n{joined}".strip() if prev else joined - ) - - return total, trace diff --git a/deprecated/mapnav/nav/nav_compose.py b/deprecated/mapnav/nav/nav_compose.py deleted file mode 100644 index 80d9f9a92..000000000 --- a/deprecated/mapnav/nav/nav_compose.py +++ /dev/null @@ -1,798 +0,0 @@ -"""Nav COMPOSE packing: parent-scoped rank + one-level indent tree. - -Parent sections are path headers only (not competing evidence units). -Child score = own_unit + w_conf * collect_confidence (default w_conf=0.5). -Group order: group_priority (external rank), then max child score, then doc order. - -Over-budget trim (progressive): keep group order; drop from back groups forward — -first non-explicit COLLECT owners, then lowest unit-score — until under budget. -""" -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple - -from ._compat import Chunk -from ._compat import line_node_id -from ._compat import ToolSpace -from .nav_address import NavLevel, address_level, owner_document - -from .nav_types import NavConfig, NavState - - -@dataclass -class ComposeFillResult: - kept_chunks: List[Chunk] - evidence_text: str - evidence_chars_actual: int - n_chunks_kept: int - truncated_last: bool - scored_chunks: List[Tuple[Chunk, float]] - dropped_any: bool = False - - -def _clamp01(x: float) -> float: - return max(0.0, min(1.0, float(x))) - - -def _section_doc_id(ts: ToolSpace, section_id: str, fallback_doc_id: str = "") -> str: - """Resolve owning document_id (registry first; never parse id strings).""" - sid = str(section_id or "").strip() - if not sid: - return str(fallback_doc_id or "") - resolved = owner_document(ts, sid, "") - if resolved: - return resolved - return str(fallback_doc_id or "") - - -def evidence_owner_section_id(chunk: Chunk) -> str: - """Section that owns a hydrated evidence chunk (strip __path/__intro suffixes).""" - nid = str(getattr(chunk, "node_id", "") or "").strip() - for suf in ("__path", "__intro", "__outline", "__self"): - if nid.endswith(suf): - return nid[: -len(suf)] - sid = str(getattr(chunk, "section_id", "") or "").strip() - return nid or sid - - -def unit_score_for_evidence_chunk(chunk: Chunk, unit_scores: Dict[str, float]) -> float: - """Map leaf/path materialize ids onto hybrid unit scores.""" - scores = unit_scores or {} - nid = str(getattr(chunk, "node_id", "") or "") - if nid.endswith("__path"): - base = nid[: -len("__path")] - return float(scores.get(base, 0.0) or 0.0) - if nid.endswith("__intro"): - base = nid[: -len("__intro")] - return float(scores.get(f"{base}__self", scores.get(base, 0.0)) or 0.0) - if nid.endswith("__self"): - base = nid[: -len("__self")] - return float(scores.get(f"{base}__self", scores.get(base, 0.0)) or 0.0) - if nid.endswith("__outline"): - base = nid[: -len("__outline")] - sid = str(getattr(chunk, "section_id", "") or "") - return float(scores.get(sid, scores.get(base, 0.0)) or 0.0) - return float(scores.get(nid, 0.0) or 0.0) - - -def direct_parent_id(ts: ToolSpace, section_id: str, doc_id: str) -> Optional[str]: - sid = str(section_id or "").strip() - if not sid: - return None - resolved = _section_doc_id(ts, sid, doc_id) - if not resolved: - return None - idx = getattr(ts, "_idx", None) - if idx is None: - return None - loc = getattr(idx, "_node_to_doc_line", {}).get(sid) - if not loc or loc[0] != resolved: - return None - _, j = loc - parents = getattr(idx, "_doc_parents", {}).get(resolved, []) - b = getattr(idx, "_bundles", {}).get(resolved) - if not b or j >= len(parents): - return None - p = parents[j] - if p is None or p < 0 or p >= len(b.lines): - return None - return line_node_id(resolved, b.lines[p].line_id) - - -def _section_title(ts: ToolSpace, section_id: str, doc_id: str) -> str: - sid = str(section_id or "").strip() - if not sid: - return "" - - # ``path_titles`` is a lightweight title lookup. Prefer it over - # ``get_structure`` because the latter may calculate the complete subtree - # chunk count, which is unnecessary while rendering evidence headers. - path_titles = getattr(ts, "path_titles", None) - if callable(path_titles): - try: - path = str(path_titles(sid, doc_id) or "").strip() - except TypeError: - path = str(path_titles(sid) or "").strip() - if path: - return path - - # Prefer structure title (Knowhere / ProviderToolSpace); never parse ids. - try: - st = ts.get_structure(sid) - except Exception: - st = None - if isinstance(st, dict): - raw = st.get("preview") or st.get("title") or "" - if isinstance(raw, str) and raw.strip(): - return raw.strip() - - resolved = _section_doc_id(ts, sid, doc_id) - idx = getattr(ts, "_idx", None) - if idx is None: - return sid - loc = getattr(idx, "_node_to_doc_line", {}).get(sid) - b = getattr(idx, "_bundles", {}).get(resolved) if loc and resolved and loc[0] == resolved else None - if not b: - level = address_level(ts, sid) - if level == NavLevel.DOCUMENT and resolved: - bb = getattr(idx, "_bundles", {}).get(resolved) - if bb and bb.lines: - title = (bb.lines[0].content or "").strip() - return title if title else sid - return sid - _, j = loc - if j < 0 or j >= len(b.lines): - return sid - title = (b.lines[j].content or "").strip() - return title if title else sid - - -def _chunk_body(chunk: Chunk) -> str: - """Strip legacy full-path [§ ...] headers; body only.""" - text = (chunk.text or "").strip() - if not text: - return "" - lines = text.splitlines() - if lines and lines[0].strip().startswith("[§"): - return "\n".join(lines[1:]).strip() - return text - - -def _line_key(chunk: Chunk) -> Tuple[int, str]: - lids = list(chunk.line_ids or ()) - if lids: - return (min(lids), chunk.node_id) - return (10**9, chunk.node_id) - - -def _child_final_score( - chunk: Chunk, - state: NavState, - *, - w_conf: float, -) -> float: - owner = evidence_owner_section_id(chunk) - own_unit = unit_score_for_evidence_chunk(chunk, dict(state.unit_scores or {})) - conf = _clamp01(float((state.collect_confidence or {}).get(owner, 0.0) or 0.0)) - return float(own_unit) + float(w_conf) * conf - - -def _parent_id_for_compose(ts: ToolSpace, section_id: str, doc_id: str) -> Optional[str]: - """Resolve a direct parent without materializing a section subtree.""" - parent_fn = getattr(ts, "parent_id", None) - if callable(parent_fn): - try: - parent = parent_fn(section_id) - except Exception: - parent = None - if parent: - return str(parent) - return direct_parent_id(ts, section_id, doc_id) - - -def _header_only_owners( - owners: set[str], ts: ToolSpace, doc_id: str -) -> set[str]: - """Find collected owners that are ancestors of another collected owner. - - Walk each owner's parent chain once instead of comparing every owner with - every other owner. This keeps compose grouping proportional to the number - of owners and hierarchy depth, even when a broad collect returns thousands - of chunks. - """ - headers: set[str] = set() - for owner in owners: - if not owner: - continue - owner_doc = _section_doc_id(ts, owner, doc_id) - current = owner - for _ in range(64): - parent = _parent_id_for_compose(ts, current, owner_doc) - if parent is None: - break - if parent in owners and _section_doc_id(ts, parent, doc_id) == owner_doc: - headers.add(parent) - current = parent - return headers - - -@dataclass -class _ChildItem: - chunk: Chunk - owner: str - score: float - line_key: Tuple[int, str] - - -@dataclass -class _ParentGroup: - parent_id: Optional[str] - parent_title: str - children: List[_ChildItem] - priority: float = 0.0 - - @property - def group_key(self) -> float: - if not self.children: - return 0.0 - return max(c.score for c in self.children) - - @property - def doc_order_key(self) -> Tuple[int, str]: - if not self.children: - return (10**9, "") - return min(c.line_key for c in self.children) - - - -def dedupe_scored(scored: Sequence[Tuple[Chunk, float]]) -> List[Tuple[Chunk, float]]: - """Keep highest score per chunk.node_id; sort by score descending.""" - best: Dict[str, Tuple[Chunk, float]] = {} - for c, score in scored: - nid = str(getattr(c, "node_id", "") or "") - if not nid: - continue - prev = best.get(nid) - if prev is None or float(score) > float(prev[1]): - best[nid] = (c, float(score)) - out = list(best.values()) - out.sort(key=lambda x: -x[1]) - return out - - -def _build_groups( - scored: Sequence[Tuple[Chunk, float]], - ts: ToolSpace, - state: NavState, - config: NavConfig, -) -> List[_ParentGroup]: - w_conf = float(getattr(config, "compose_confidence_weight", 0.5) or 0.5) - seen_ids: set[str] = set() - items: List[Tuple[Chunk, str, float]] = [] - for chunk, _bag in scored: - nid = str(getattr(chunk, "node_id", "") or "") - if not nid or nid in seen_ids: - continue - body = _chunk_body(chunk) - if not body: - continue - seen_ids.add(nid) - owner = evidence_owner_section_id(chunk) - score = _child_final_score(chunk, state, w_conf=w_conf) - items.append((chunk, owner, score)) - - owners = {owner for _c, owner, _s in items if owner} - header_owners = _header_only_owners(owners, ts, state.doc_id) - - groups: Dict[Optional[str], _ParentGroup] = {} - for chunk, owner, score in items: - if owner in header_owners: - continue - owner_doc = _section_doc_id(ts, owner, state.doc_id) or str( - getattr(chunk, "doc_id", "") or "" - ) - parent_id = direct_parent_id(ts, owner, owner_doc) - if parent_id is None: - parent_id = owner - if parent_id not in groups: - title = _section_title(ts, parent_id, owner_doc) - groups[parent_id] = _ParentGroup( - parent_id=parent_id, - parent_title=title, - children=[], - priority=float((state.group_priority or {}).get(parent_id, 0.0) or 0.0), - ) - groups[parent_id].children.append( - _ChildItem( - chunk=chunk, - owner=owner, - score=score, - line_key=_line_key(chunk), - ) - ) - return list(groups.values()) - - -def _render_group( - group: _ParentGroup, - selected: Sequence[_ChildItem], - *, - evidence_index: int, -) -> str: - """Render one evidence block as ``[E#]`` + path + bodies.""" - bodies = [_chunk_body(child.chunk) for child in selected] - return _render_evidence_blocks( - [(group.parent_title or "", bodies)], - start_index=evidence_index, - ) - - -def _render_evidence_blocks( - groups: Sequence[tuple[str, Sequence[str]]], - *, - start_index: int = 1, -) -> str: - parts: list[str] = [] - index = max(1, int(start_index or 1)) - for path, bodies in groups: - texts = [str(t or "").strip() for t in bodies] - texts = [t for t in texts if t] - if not texts: - continue - block: list[str] = [f"[E{index + len(parts)}]"] - header = str(path or "").strip() - if header: - block.append(f"[§ {header}]") - indent = len(texts) >= 2 - for text in texts: - if indent: - block.append( - "\n".join( - (" " + ln if ln.strip() else ln) for ln in text.splitlines() - ) - ) - else: - block.append(text) - parts.append("\n".join(block).strip()) - return "\n\n".join(parts) - - -def _scored_flat(groups: Sequence[_ParentGroup]) -> List[Tuple[Chunk, float]]: - out: List[Tuple[Chunk, float]] = [] - for g in groups: - for c in sorted(g.children, key=lambda x: (-x.score, x.line_key)): - out.append((c.chunk, c.score)) - return out - - -def _n_pool_children(groups: Sequence[_ParentGroup]) -> int: - return sum(len(g.children) for g in groups) - - -def _is_explicit_collect(child: _ChildItem, state: NavState) -> bool: - """True if LLM explicitly COLLECTed this owner (or left non-zero confidence).""" - owner = str(child.owner or "").strip() - if not owner: - return False - explicit = getattr(state, "explicit_collect_ids", None) or set() - if owner in explicit: - return True - conf = float((state.collect_confidence or {}).get(owner, 0.0) or 0.0) - return conf > 0.0 - - -def _unit_score(child: _ChildItem, state: NavState) -> float: - return float( - unit_score_for_evidence_chunk(child.chunk, dict(state.unit_scores or {})) - ) - - -def _render_kept( - groups: Sequence[_ParentGroup], - kept_ids: Set[str], - *, - budget_chars: int, - min_partial_chars: int = 20, -) -> ComposeFillResult: - """Render kept children in group order (doc order within group), full text.""" - scored_flat = _scored_flat(groups) - parts: List[str] = [] - kept_chunks: List[Chunk] = [] - used = 0 - truncated_last = False - sep = "\n\n" - n_pool = _n_pool_children(groups) - - for g in groups: - entries = [ - c - for c in sorted(g.children, key=lambda x: x.line_key) - if c.chunk.node_id in kept_ids - ] - if not entries: - continue - block = _render_group( - g, entries, evidence_index=len(parts) + 1 - ) - add = len(block) + (len(sep) if parts else 0) - if used + add <= budget_chars: - parts.append(block) - kept_chunks.extend(c.chunk for c in entries) - used += add - continue - remain = budget_chars - used - (len(sep) if parts else 0) - if remain >= min_partial_chars and not parts: - parts.append(block[:budget_chars]) - kept_chunks.extend(c.chunk for c in entries) - used = budget_chars - truncated_last = True - break - if remain >= min_partial_chars: - parts.append(block[:remain]) - kept_chunks.extend(c.chunk for c in entries) - used = budget_chars - truncated_last = True - break - - text = sep.join(parts) - return ComposeFillResult( - kept_chunks=kept_chunks, - evidence_text=text, - evidence_chars_actual=len(text), - n_chunks_kept=len(kept_chunks), - truncated_last=truncated_last, - scored_chunks=scored_flat, - dropped_any=len(kept_chunks) < n_pool or truncated_last, - ) - - -def _count_power10_boundaries(lower: int, upper: int) -> int: - """Count decimal digit-boundary positions in ``(lower, upper]``.""" - if upper <= lower: - return 0 - boundary = 10 - count = 0 - while boundary <= upper: - if boundary > lower: - count += 1 - boundary *= 10 - return count - - -class _SelectionLength: - """Maintain exact rendered selection length under one-child mutations.""" - - def __init__(self, groups: Sequence[_ParentGroup], kept_ids: Set[str]) -> None: - self._groups = groups - self.kept_ids = kept_ids - self._group_by_node_id: Dict[str, int] = {} - self._plain_lengths: Dict[str, int] = {} - self._indented_lengths: Dict[str, int] = {} - self._counts: List[int] = [] - self._plain_totals: List[int] = [] - self._indented_totals: List[int] = [] - self._nonempty_count = 0 - self.total = 0 - - for group_index, group in enumerate(groups): - plain_total = 0 - indented_total = 0 - count = 0 - for child in group.children: - node_id = str(child.chunk.node_id or "") - if not node_id: - continue - body = _chunk_body(child.chunk) - plain_length = len(body) - indented_length = len( - "\n".join( - (" " + line if line.strip() else line) - for line in body.splitlines() - ) - ) - self._group_by_node_id[node_id] = group_index - self._plain_lengths[node_id] = plain_length - self._indented_lengths[node_id] = indented_length - if node_id in kept_ids: - count += 1 - plain_total += plain_length - indented_total += indented_length - self._counts.append(count) - self._plain_totals.append(plain_total) - self._indented_totals.append(indented_total) - - self._recalculate_total() - - @staticmethod - def _block_length( - group: _ParentGroup, - *, - count: int, - plain_total: int, - indented_total: int, - evidence_index: int, - ) -> int: - if count <= 0: - return 0 - prefix = len(f"[E{evidence_index}]") - if group.parent_title: - prefix += 1 + len(f"[§ {group.parent_title}]") - body_length = indented_total if count >= 2 else plain_total - return prefix + count + body_length - - def _recalculate_total(self) -> None: - total = 0 - evidence_index = 0 - for index, group in enumerate(self._groups): - count = self._counts[index] - if count <= 0: - continue - evidence_index += 1 - total += self._block_length( - group, - count=count, - plain_total=self._plain_totals[index], - indented_total=self._indented_totals[index], - evidence_index=evidence_index, - ) - self._nonempty_count = evidence_index - self.total = total + max(0, evidence_index - 1) * 2 - - def _evidence_index(self, group_index: int) -> int: - return sum(1 for count in self._counts[:group_index] if count > 0) + 1 - - def remove(self, node_id: str) -> None: - group_index = self._group_by_node_id.get(node_id) - if group_index is None or node_id not in self.kept_ids: - return - old_count = self._counts[group_index] - old_index = self._evidence_index(group_index) - old_block = self._block_length( - self._groups[group_index], - count=old_count, - plain_total=self._plain_totals[group_index], - indented_total=self._indented_totals[group_index], - evidence_index=old_index, - ) - self.kept_ids.remove(node_id) - self._counts[group_index] -= 1 - self._plain_totals[group_index] -= self._plain_lengths[node_id] - self._indented_totals[group_index] -= self._indented_lengths[node_id] - if old_count > 1: - new_block = self._block_length( - self._groups[group_index], - count=old_count - 1, - plain_total=self._plain_totals[group_index], - indented_total=self._indented_totals[group_index], - evidence_index=old_index, - ) - self.total += new_block - old_block - return - - self.total -= old_block - if self._nonempty_count > 1: - self.total -= 2 - self.total -= _count_power10_boundaries(old_index, self._nonempty_count) - self._nonempty_count -= 1 - - def add(self, node_id: str) -> bool: - group_index = self._group_by_node_id.get(node_id) - if group_index is None or node_id in self.kept_ids: - return False - old_count = self._counts[group_index] - evidence_index = self._evidence_index(group_index) - if old_count > 0: - old_block = self._block_length( - self._groups[group_index], - count=old_count, - plain_total=self._plain_totals[group_index], - indented_total=self._indented_totals[group_index], - evidence_index=evidence_index, - ) - - self.kept_ids.add(node_id) - self._counts[group_index] += 1 - self._plain_totals[group_index] += self._plain_lengths[node_id] - self._indented_totals[group_index] += self._indented_lengths[node_id] - new_block = self._block_length( - self._groups[group_index], - count=self._counts[group_index], - plain_total=self._plain_totals[group_index], - indented_total=self._indented_totals[group_index], - evidence_index=evidence_index, - ) - if old_count > 0: - self.total += new_block - old_block - else: - old_nonempty_count = self._nonempty_count - self.total += new_block - if old_nonempty_count > 0: - self.total += 2 - self.total += _count_power10_boundaries( - evidence_index, old_nonempty_count + 1 - ) - self._nonempty_count += 1 - return True - - -def _selection_chars( - groups: Sequence[_ParentGroup], kept_ids: Set[str] -) -> int: - """Exact rendered length of the current full-text selection.""" - return _SelectionLength(groups, kept_ids).total - - -def _refill_to_budget( - groups: Sequence[_ParentGroup], - kept_ids: Set[str], - state: NavState, - *, - budget_chars: int, -) -> None: - """Re-add dropped children, best first, while the rendered selection still fits. - - Dropping runs back-to-front and stops the moment the selection fits, so whole - tail groups can be gone while a large slice of the budget sits unused. Without - this pass a single oversized chunk near the cut point strands everything behind - it (observed: 4 chunks kept, 5.7k of 12k budget idle, 414 dropped chunks that - each still fit). - """ - dropped = [ - c - for g in groups - for c in g.children - if c.chunk.node_id and c.chunk.node_id not in kept_ids - ] - if not dropped: - return - dropped.sort( - key=lambda c: ( - not _is_explicit_collect(c, state), - -_unit_score(c, state), - c.line_key, - ) - ) - selection = _SelectionLength(groups, kept_ids) - for child in dropped: - if selection.total >= budget_chars: - break - if len(child.chunk.text or "") > budget_chars - selection.total: - continue - if not selection.add(child.chunk.node_id): - continue - if selection.total > budget_chars: - selection.remove(child.chunk.node_id) - - -def _pack_trim( - groups: List[_ParentGroup], - state: NavState, - *, - budget_chars: int, - min_partial_chars: int = 20, -) -> ComposeFillResult: - """Keep group order; drop back-to-front until under budget, then refill slack. - - Phase 1: drop non-explicit COLLECT owners (later groups first, later doc first). - Phase 2: drop remaining by lowest unit score (later groups first). - Phase 3: re-add dropped children best-first while they still fit. - """ - scored_flat = _scored_flat(groups) - if budget_chars <= 0 or not groups: - return ComposeFillResult([], "", 0, 0, False, scored_flat, dropped_any=False) - - kept_ids: Set[str] = { - c.chunk.node_id for g in groups for c in g.children if c.chunk.node_id - } - if not kept_ids: - return ComposeFillResult([], "", 0, 0, False, scored_flat, dropped_any=False) - - selection = _SelectionLength(groups, kept_ids) - - def fits() -> bool: - return selection.total <= budget_chars - - if fits(): - return _render_kept( - groups, kept_ids, budget_chars=budget_chars, min_partial_chars=min_partial_chars - ) - - # Phase 1: non-explicit, back group → front; within group, later doc first. - for g in reversed(groups): - candidates = [ - c - for c in g.children - if c.chunk.node_id in kept_ids and not _is_explicit_collect(c, state) - ] - candidates.sort(key=lambda c: c.line_key, reverse=True) - for child in candidates: - if fits(): - break - selection.remove(child.chunk.node_id) - if fits(): - break - - # Phase 2: still over → lowest unit first, back group → front. - if not fits(): - for g in reversed(groups): - candidates = [c for c in g.children if c.chunk.node_id in kept_ids] - candidates.sort( - key=lambda c: (_unit_score(c, state), c.line_key) - ) - for child in candidates: - if fits(): - break - if len(kept_ids) <= 1: - break - selection.remove(child.chunk.node_id) - if fits() or len(kept_ids) <= 1: - break - - _refill_to_budget(groups, kept_ids, state, budget_chars=budget_chars) - - return _render_kept( - groups, kept_ids, budget_chars=budget_chars, min_partial_chars=min_partial_chars - ) - - - -def pack_nav_evidence( - collected: Sequence[Tuple[Chunk, float]], - ts: ToolSpace, - state: NavState, - config: NavConfig, - *, - budget_chars: int, - min_partial_chars: int = 20, -) -> ComposeFillResult: - """Pack collected evidence under parent scopes into a budgeted tree string. - - Only progressive trim+refill (``_pack_trim``). Waterfill/greedy paths removed. - """ - if budget_chars <= 0: - return ComposeFillResult([], "", 0, 0, False, [], dropped_any=False) - - groups = _build_groups(collected, ts, state, config) - groups.sort(key=lambda g: (-g.priority, -g.group_key, g.doc_order_key)) - return _pack_trim( - groups, - state, - budget_chars=budget_chars, - min_partial_chars=min_partial_chars, - ) - - -def parse_collect_confidence( - obj: Dict[str, Any], - selected: Sequence[Any], -) -> Dict[str, float]: - """Map selected LegalActions -> confidence in [0,1]. Missing => 0.""" - raw = (obj or {}).get("confidence") - out: Dict[str, float] = {} - if isinstance(raw, (int, float)): - c = _clamp01(float(raw)) - for a in selected: - aid = str(getattr(a, "action_id", "") or "").strip().upper() - if aid: - out[aid] = c - return out - if isinstance(raw, dict): - norm = {str(k).strip().upper(): v for k, v in raw.items()} - for a in selected: - aid = str(getattr(a, "action_id", "") or "").strip().upper() - if not aid: - continue - if aid in norm and norm[aid] is not None: - try: - out[aid] = _clamp01(float(norm[aid])) - except (TypeError, ValueError): - out[aid] = 0.0 - else: - out[aid] = 0.0 - return out - for a in selected: - aid = str(getattr(a, "action_id", "") or "").strip().upper() - if aid: - out[aid] = 0.0 - return out diff --git a/deprecated/mapnav/nav/nav_control.py b/deprecated/mapnav/nav/nav_control.py deleted file mode 100644 index 698fa5874..000000000 --- a/deprecated/mapnav/nav/nav_control.py +++ /dev/null @@ -1,371 +0,0 @@ -"""Merged check authority (PLAN×NAV fusion): one planner decision per wave. - -``plan_control`` is the single authority that reconciles each wave's new -evidence against the plan's coverage checklist. It returns one decision per -subgoal plus one global decision. ``REPLAN`` can only originate here. -""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from typing import Any, Dict, List, Literal, Optional, Sequence - -from .nav_plan import RetrievalPlan, Subgoal -from .nav_policy import _extract_json_obj -from .nav_types import NavConfig, NavState - -SubgoalDecisionKind = Literal["accept", "widen", "drop"] -GlobalDecisionKind = Literal["continue", "replan", "done"] - -_CONTROL_PURPOSE = "nav_plan_control_v2" -_SUBGOAL_DECISIONS = {"accept", "widen", "drop"} -_GLOBAL_DECISIONS = {"continue", "replan", "done"} - - -@dataclass -class SubgoalDecision: - subgoal_id: str - decision: SubgoalDecisionKind = "accept" - note: str = "" - - -@dataclass -class PlanControlDecision: - per_subgoal: Dict[str, SubgoalDecision] = field(default_factory=dict) - global_action: GlobalDecisionKind = "continue" - reason: str = "" - raw: str = "" - - -def _section_summary_text(ts: Any, section_id: str) -> str: - """Full stored section summary (already head/tail clipped at build time).""" - sid = str(section_id or "").strip() - if not sid: - return "" - try: - from .nav_address import owner_document - from section_summary_store import get_summary - - doc = owner_document(ts, sid, "") if ts is not None else "" - text = str(get_summary(sid, doc_id=doc) or "").strip() - if text: - return text - except Exception: - pass - if ts is not None: - try: - st = ts.get_structure(sid) or {} - text = str(st.get("summary") or "").strip() - if text: - return text - # Last resort: title/preview only (still no raw-body truncation here). - return str(st.get("preview") or "").strip() - except Exception: - return "" - return "" - - -def _digest_collected_summaries( - ts: Any, - *, - collected_section_ids: Optional[Sequence[str]] = None, -) -> str: - """One line per explicit COLLECT target: path + its prebuilt summary. - - Hydration descendants are intentionally omitted — control only needs the - nodes the harvester chose, not every leaf sucked in under them. - """ - sids = [str(s).strip() for s in (collected_section_ids or []) if str(s).strip()] - if not sids: - # No explicit COLLECTs this wave: do not invent digest lines from - # hydrated chunk owners (that reintroduces the descendant flood). - return "" - lines: List[str] = [] - for sid in sids: - title = _section_display_title(ts, sid) - summary = _section_summary_text(ts, sid) - if summary: - lines.append(f"- {title}: {summary}") - else: - lines.append(f"- {title}: (no summary)") - return "\n".join(lines) - - -def _section_display_title(ts: Any, section_id: str) -> str: - """Human title for control digest — never parse section_id strings.""" - sid = str(section_id or "").strip() - if not sid: - return "" - if ts is not None: - try: - st = ts.get_structure(sid) or {} - title = str(st.get("preview") or st.get("title") or "").strip() - if title: - return title - except Exception: - pass - return sid - - -def _wave_subgoal_block( - subgoal: Subgoal, - *, - digest: str, - harvest_meta: Optional[Dict[str, Any]], - attempt_count: int, -) -> str: - card = subgoal.contract.cardinality - contract_line = f"{subgoal.contract.kind}" + ( - f" cardinality={card}" if card is not None else "" - ) - lines = [ - f"[{subgoal.id}] need: {subgoal.need}", - f" contract: {contract_line}", - " search_space: shared (no per-subgoal scope/anchor)", - f" attempt: {attempt_count}", - ] - if harvest_meta: - lines.append( - " harvest: " - f"visited={len(harvest_meta.get('visited_section_ids') or [])} " - f"policy_calls={harvest_meta.get('n_policy_calls', 0)}" - ) - harvest_reason = str(harvest_meta.get("reason") or "").strip() - if harvest_reason: - lines.append(f" harvest_reason: {harvest_reason}") - lines.append(f" explicit_collect_summaries:\n{digest.strip() or ' (empty)'}") - return "\n".join(lines) - - -def _plan_overview_block(plan: RetrievalPlan, state: NavState) -> str: - lines = ["coverage_checklist:"] - checklist = list(getattr(plan, "coverage_checklist", None) or []) - if checklist: - for item in checklist: - lines.append(f"- {item.id}: {item.fact}") - else: - lines.append("- (none)") - lines.append("subgoals:") - for sg in plan.subgoals: - status = ( - "satisfied" - if sg.id in state.satisfied_subgoal_ids - else "attempted" if sg.id in state.attempted_subgoal_ids else "pending" - ) - lines.append(f"- {sg.id}: {sg.need} [{status}]") - return "\n".join(lines) - - -def _control_system_prompt() -> str: - return ( - "You are the single retrieval-plan controller for one wave of a " - "hierarchical document retrieval episode. You replace all separate " - "per-region stop/retry judgments with one decision.\n\n" - "The plan has a global coverage_checklist (facts that must appear in " - "episode evidence) and one shared search space. For each subgoal in " - "this wave you are shown: its need/contract, attempt count, " - "the harvester's own explanation (harvest_reason), and the section " - "summaries for nodes explicitly COLLECTed THIS wave only (never " - "hydration descendants, and never older evidence from other subgoals; " - "summaries are the prebuilt node summaries).\n\n" - "=== Per-subgoal decisions ===\n" - " - accept: this subgoal's need is covered by this wave's evidence.\n" - " - widen: not yet covered; keep the subgoal unsettled for another " - "harvest from the shared root. Prior dead-ends stay dismissed; PLAN " - "rewrites this subgoal's retrieval_query from your note. " - "Do not name an entry point.\n" - " - drop: not covered and further attempts are unlikely to help " - "(e.g. evidence is structurally absent, or harvest_reason shows the " - "search has already reached the document root with nothing found); " - "stop trying this subgoal.\n\n" - "=== Global decision ===\n" - " - continue: proceed to the next wave with current subgoal set.\n" - " - replan: the retrieval plan's decomposition itself is wrong " - "(e.g. missing checklist facts, wrong dependencies) and needs regenerating. " - "Only choose this for a structural plan problem, not for an " - "individual subgoal's evidence gap (use widen/drop for those).\n" - " - done: the coverage_checklist is met (or cannot be met " - "further); stop the episode.\n\n" - "Return ONLY one JSON object:\n" - "{\n" - ' "subgoals": {"s1": {"decision": "accept|widen|drop", ' - '"note": "..."}},\n' - ' "global": "continue|replan|done",\n' - ' "reason": "..."\n' - "}\n" - "Do not include any explanation outside the JSON.\n\n" - "IMPORTANT:\n" - "1. All agent-generated text (note/reason) MUST be in English.\n" - "2. Keep reason under 30 words; keep each note under 15 words.\n" - ) - - -def _parse_control_decision(obj: Dict[str, Any], subgoal_ids: Sequence[str]) -> PlanControlDecision: - per_subgoal: Dict[str, SubgoalDecision] = {} - raw_sub = obj.get("subgoals") if isinstance(obj, dict) else None - if isinstance(raw_sub, dict): - for sid in subgoal_ids: - row = raw_sub.get(sid) - if not isinstance(row, dict): - continue - decision = str(row.get("decision") or "").strip().lower() - if decision not in _SUBGOAL_DECISIONS: - decision = "accept" - per_subgoal[sid] = SubgoalDecision( - subgoal_id=sid, - decision=decision, # type: ignore[arg-type] - note=str(row.get("note") or "")[:200], - ) - global_action = str(obj.get("global") or "").strip().lower() - if global_action not in _GLOBAL_DECISIONS: - global_action = "continue" - return PlanControlDecision( - per_subgoal=per_subgoal, - global_action=global_action, # type: ignore[arg-type] - reason=str(obj.get("reason") or "")[:300], - ) - - -def _fallback_decision( - subgoal_ids: Sequence[str], - signals: Dict[str, Any], -) -> PlanControlDecision: - """Deterministic fallback when the LLM call fails or returns malformed JSON. - - Only mechanical signal left: this wave collected evidence → accept, - otherwise widen. Checklist judgment is not attempted offline. - """ - per_subgoal: Dict[str, SubgoalDecision] = {} - for sid in subgoal_ids: - sig = signals.get(sid) - has_evidence = ( - int(getattr(sig, "chars_used", 0) or 0) > 0 if sig is not None else False - ) - per_subgoal[sid] = SubgoalDecision( - subgoal_id=sid, - decision="accept" if has_evidence else "widen", - ) - return PlanControlDecision(per_subgoal=per_subgoal, global_action="continue", reason="fallback") - - -def plan_control( - ts: Any, - state: NavState, - config: NavConfig, - *, - plan: RetrievalPlan, - wave_outputs: Sequence[Dict[str, Any]], -) -> PlanControlDecision: - """One LLM call per wave: per-subgoal accept/widen/drop + global signal.""" - from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted - - by_id = {s.id: s for s in plan.subgoals} - subgoal_ids = [str(item.get("subgoal_id")) for item in wave_outputs] - signals = {str(item.get("subgoal_id")): item.get("result") for item in wave_outputs} - if not subgoal_ids: - return PlanControlDecision(global_action="continue") - if nav_token_budget_exhausted(): - decision = _fallback_decision(subgoal_ids, signals) - decision.global_action = "done" - decision.reason = "token_limit" - return decision - - blocks: List[str] = [] - for item in wave_outputs: - sid = str(item.get("subgoal_id")) - sg = by_id.get(sid) - if sg is None: - continue - result = item.get("result") - # Control digest = explicit COLLECT targets only (not hydrated descendants). - if hasattr(result, "explicit_collect_ids"): - digest_ids = list(getattr(result, "explicit_collect_ids", None) or []) - else: - digest_ids = list(getattr(result, "collected_section_ids", None) or []) - digest = _digest_collected_summaries( - ts, - collected_section_ids=digest_ids, - ) - blocks.append( - _wave_subgoal_block( - sg, - digest=digest, - harvest_meta=item.get("harvest"), - attempt_count=int(state.subgoal_attempt_counts.get(sid, 0)), - ) - ) - user = ( - f"=== Plan Overview ===\n{_plan_overview_block(plan, state)}\n" - "=== End Plan Overview ===\n\n" - f"=== This Wave ({len(wave_outputs)} subgoal(s)) ===\n" - + "\n\n".join(blocks) - + "\n=== End This Wave ===\n\n" - "Return the control decision JSON." - ) - - try: - from .nav_llm import ( # type: ignore - nav_chat, - resolve_nav_model, - ) - - model = resolve_nav_model( - model=config.planner_model, - model_env="NAV_PLANNER_MODEL", - fallback_envs=("NAV_LLM_MODEL",), - ) - # Control is short JSON (accept/widen/drop); never use planner thinking. - max_tokens = max( - 256, - int(getattr(config, "planner_llm_max_tokens", 0) or 0) - or int(config.llm_max_tokens or 256), - ) - timeout_s = float(os.environ.get("NAV_PLANNER_TIMEOUT_SECONDS", "").strip() or "0") - if timeout_s <= 0: - timeout_s = 90.0 - cached = nav_chat( - purpose=_CONTROL_PURPOSE, - model=model, - messages=[ - {"role": "system", "content": _control_system_prompt()}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=max_tokens, - response_format={"type": "json_object"}, - thinking_role="action", - context="Nav Plan Control", - api_key_env="NAV_PLANNER_API_KEY", - base_url_env="NAV_PLANNER_BASE_URL", - timeout=timeout_s, - usage_tag="nav_plan_control", - ) - text = str(cached.get("content") or "").strip() - if not text: - reasoning = str(cached.get("reasoning_content") or "").strip() - if reasoning and _extract_json_obj(reasoning): - text = reasoning - obj = _extract_json_obj(text) or {} - decision = _parse_control_decision(obj, subgoal_ids) - decision.raw = text[:1000] - # Any subgoal the model omitted still needs an explicit decision. - for sid in subgoal_ids: - if sid not in decision.per_subgoal: - sig = signals.get(sid) - has_evidence = ( - int(getattr(sig, "chars_used", 0) or 0) > 0 - if sig is not None - else False - ) - decision.per_subgoal[sid] = SubgoalDecision( - subgoal_id=sid, decision="accept" if has_evidence else "widen" - ) - return decision - except NavTokenLimit: - decision = _fallback_decision(subgoal_ids, signals) - decision.global_action = "done" - decision.reason = "token_limit" - return decision - except Exception: - return _fallback_decision(subgoal_ids, signals) diff --git a/deprecated/mapnav/nav/nav_harvest.py b/deprecated/mapnav/nav/nav_harvest.py deleted file mode 100644 index 9fc894e38..000000000 --- a/deprecated/mapnav/nav/nav_harvest.py +++ /dev/null @@ -1,480 +0,0 @@ -"""One-shot harvest primitive (PLAN×NAV fusion). - -Replaces the multi-step COLLECT/DISPATCH/FINISH ReAct loop with a single -policy decision per visible map node: the model returns the union of nodes to -collect and the union of nodes to dispatch to a deeper harvester in one call. -FINISH is implicit — selecting neither collect nor dispatch simply ends this -region; there is no separate round trip asking "are you done?". - -Recursion is bounded by ``max_harvest_depth`` and only ever follows an -explicit DISPATCH selection (never a fixed step budget). The map already -folds large scopes to title-only and offers D* ids for internal nodes -(``nav_projection.build_map`` / ``nav_actions.build_legal_actions``), so a -node whose children overflow the display budget naturally nudges the model -toward DISPATCH — no separate "scope overflow" special case is needed here. - -Depends only on the existing map/action primitives (``nav_projection``, -``nav_actions``) plus ``nav_navigate._apply_collect`` for hydration. No new -ToolSpace capability is required beyond the 5 documented in -docs/audit_plan_nav_overlap.md. -""" - -from __future__ import annotations - -import re -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple - -from .nav_actions import build_legal_actions, format_actionable_map_observation -from .nav_compose import parse_collect_confidence -from .nav_plan import Subgoal -from .nav_policy import _extract_json_obj # reuse: same tolerant JSON extraction -from .nav_projection import build_map -from .nav_types import ActionKind, LegalAction, NavConfig, NavState, Projection - -_HARVEST_PURPOSE_DEPTH0 = "nav_harvest_v1" -_HARVEST_PURPOSE_CHILD = "nav_harvest_child_v1" - - -def _extract_id_list_field(text: str, field: str) -> List[str]: - """Pull a JSON string-array field even when the surrounding object is truncated.""" - m = re.search( - rf'"{field}"\s*:\s*\[(.*?)\]', - text or "", - flags=re.S | re.I, - ) - if not m: - return [] - return [ - str(x).strip().upper() - for x in re.findall(r'"([^"]+)"', m.group(1)) - if str(x).strip() - ] - - -@dataclass -class HarvestResult: - subgoal_id: str - new_section_ids: List[str] = field(default_factory=list) - visited_section_ids: List[str] = field(default_factory=list) - n_policy_calls: int = 0 - max_depth_hit: bool = False - # ": " per policy call this harvest tree made, joined by - # " | " — plan_control reads this alongside collected section summaries to judge widen - # vs drop (previously only reached AgentStep.detail, never the controller). - reason: str = "" - search_assets: List[Dict[str, Any]] = field(default_factory=list) - n_assets_added: int = 0 - - -def _actions_by_ids(actions: Sequence[LegalAction], ids: Sequence[str]) -> List[LegalAction]: - from .nav_actions import actions_by_ids - - return actions_by_ids(list(actions), list(ids)) - - -def _harvest_system_prompt(*, dispatch_available: bool) -> str: - dispatch_rule = ( - " - dispatch=D*: hand a node's subtree to a deeper harvester before " - "deciding; use this when the node's title/summary alone does not " - "already tell you whether the needed evidence is inside — the deeper " - "harvester will make its own single decision over that subtree.\n" - if dispatch_available - else "" - ) - return ( - "You are a retrieval harvester working on ONE retrieval subgoal inside " - "a hierarchical document map region.\n\n" - "In a single response, decide which visible nodes to collect as " - "evidence for the subgoal below, and (if available) which visible " - "nodes to hand to a deeper harvester. Selecting neither finishes this " - "region — there is no separate finish step.\n\n" - "=== Rules ===\n\n" - " - collect=C*: add the node (and its full subtree, when it is a " - "parent) as evidence for the subgoal's contract.\n" - f"{dispatch_rule}" - " - Use only action IDs shown on a node line. Never invent IDs or " - "write raw section paths as targets.\n" - " - Select every node that plausibly bears on the subgoal (collect " - "if the answer looks to be inside it already, dispatch if you are " - "unsure and need to look deeper). Leave a node unselected only when " - "it clearly does not relate to the subgoal.\n" - " - A node you leave unselected this call will not be shown again " - "for this subgoal, so when in doubt prefer dispatch over silently " - "skipping a node.\n" - " - Provide confidence in [0,1] for every collect id (object map " - "keyed by action id, or a single scalar when there is exactly one).\n" - " - search_assets: optional list of {kind:\"images\"|\"tables\", " - "query, scope?} when the subgoal needs figure/table evidence only " - "(or in addition). kind filters Knowhere chunk_type; scope defaults " - "to the current region and must stay inside one document " - "(document root allowed; never the whole namespace). " - "Asset search does not dismiss map nodes.\n\n" - "=== End Rules ===\n\n" - "Return ONLY one JSON object, e.g.:\n" - '{"collect_ids":["C1"],"dispatch_ids":[],"search_assets":[],' - '"confidence":{"C1":0.8},"reason":"short reason"}\n' - "Do not include any explanation outside the JSON.\n\n" - "IMPORTANT:\n" - "1. All agent-generated text (reason) MUST be in English.\n" - "2. Document content and section titles MUST remain in their " - "original language.\n" - "3. Keep reason under 25 words.\n" - ) - - -def _harvest_user_prompt( - *, - subgoal: Subgoal, - query: str, - observation: str, -) -> str: - card = subgoal.contract.cardinality - contract_line = f"{subgoal.contract.kind}" + ( - f" cardinality={card}" if card is not None else "" - ) - return ( - f"Subgoal need: {subgoal.need or query}\n" - f"Retrieval query: {query}\n" - f"Contract: {contract_line}\n\n" - f"=== Region Observation ===\n{observation}\n=== End Region Observation ===\n" - ) - - -def _rule_fallback_selection( - actions: Sequence[LegalAction], -) -> Tuple[List[LegalAction], List[LegalAction]]: - """Deterministic pick when the LLM returns no valid ids: one collect else nothing.""" - for a in actions: - if a.kind == ActionKind.COLLECT: - return [a], [] - return [], [] - - -def harvest_policy_call( - ts: Any, - state: NavState, - config: NavConfig, - *, - subgoal: Subgoal, - query: str, - projection: Projection, - actions: Sequence[LegalAction], - depth: int, -) -> Tuple[List[LegalAction], List[LegalAction], Dict[str, float], str, Dict[str, Any]]: - """One LLM call; ``meta["search_assets"]`` holds normalized asset requests.""" - from .nav_assets import parse_search_assets - from .nav_llm import nav_chat, resolve_nav_model - from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted - - if nav_token_budget_exhausted(): - return [], [], {}, "token_limit", { - "stop_reason": "token_limit", - "search_assets": [], - "depth": depth, - } - - model = resolve_nav_model( - model=(config.subagent_model if depth > 0 else config.llm_model), - model_env=("NAV_SUBAGENT_MODEL" if depth > 0 else "NAV_LLM_MODEL"), - fallback_envs=("NAV_LLM_MODEL",), - ) - - dispatch_available = any(a.kind == ActionKind.DISPATCH for a in actions) - system = _harvest_system_prompt(dispatch_available=dispatch_available) - observation = format_actionable_map_observation( - projection, list(actions), inline_summary=projection.scope is not None - ) - asset_ctx = str(getattr(state, "asset_observation_context", "") or "").strip() - if asset_ctx: - observation = f"{asset_ctx}\n\n{observation}" - user = _harvest_user_prompt(subgoal=subgoal, query=query, observation=observation) - purpose = _HARVEST_PURPOSE_DEPTH0 if depth == 0 else _HARVEST_PURPOSE_CHILD - - try: - cached = nav_chat( - purpose=purpose, - model=model, - messages=[ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=max( - 256, - int(getattr(config, "harvest_llm_max_tokens", 0) or 0) - or int(config.llm_max_tokens), - ), - response_format={"type": "json_object"}, - context="Nav Harvest", - usage_tag="nav_harvest", - ) - except NavTokenLimit: - return [], [], {}, "token_limit", { - "stop_reason": "token_limit", - "search_assets": [], - "depth": depth, - } - text = str(cached.get("content") or "").strip() - obj = _extract_json_obj(text) or {} - - collect_ids = [str(x).strip().upper() for x in (obj.get("collect_ids") or []) if str(x).strip()] - dispatch_ids = [str(x).strip().upper() for x in (obj.get("dispatch_ids") or []) if str(x).strip()] - # Truncated completions (hit llm_max_tokens mid-JSON) fail object parse and - # used to look like an intentional empty selection — recover id arrays. - if not collect_ids and not dispatch_ids and text: - collect_ids = _extract_id_list_field(text, "collect_ids") - dispatch_ids = _extract_id_list_field(text, "dispatch_ids") - search_assets = parse_search_assets(obj.get("search_assets")) - collect_actions = [a for a in _actions_by_ids(actions, collect_ids) if a.kind == ActionKind.COLLECT] - dispatch_actions = [a for a in _actions_by_ids(actions, dispatch_ids) if a.kind == ActionKind.DISPATCH] - fallback_used = False - if ( - not collect_actions - and not dispatch_actions - and not collect_ids - and not dispatch_ids - and not search_assets - ): - # Genuinely empty selection == implicit finish; not a parse failure. - pass - elif not collect_actions and not dispatch_actions and not search_assets: - collect_actions, dispatch_actions = _rule_fallback_selection(actions) - fallback_used = True - - confidence = parse_collect_confidence(obj, collect_actions) - reason = str(obj.get("reason") or "")[:300] - meta = { - "model": model, - "raw": text[:500], - "depth": depth, - "fallback_used": fallback_used, - "search_assets": search_assets, - } - return collect_actions, dispatch_actions, confidence, reason, meta - - -def harvest( - ts: Any, - state: NavState, - config: NavConfig, - *, - subgoal: Subgoal, - entry_scope: Optional[str], - query: str, - steps_out: Optional[List[Any]] = None, - allowed_section_ids: Optional[Set[str]] = None, -) -> HarvestResult: - """Single-decision-per-node evidence harvest for one subgoal.""" - result = HarvestResult(subgoal_id=subgoal.id) - from .nav_address import NavLevel, address_level - - if entry_scope is None: - initial_depth = 0 - else: - level = address_level(ts, entry_scope) - if level in (NavLevel.NAMESPACE, NavLevel.DOCUMENT): - initial_depth = 0 - else: - # SECTION / unknown → already inside a document. - initial_depth = 1 - _harvest_node( - ts, - state, - config, - subgoal=subgoal, - node_scope=entry_scope, - query=query, - depth=initial_depth, - steps_out=steps_out, - result=result, - allowed_section_ids=allowed_section_ids, - ) - return result - - -def _harvest_node( - ts: Any, - state: NavState, - config: NavConfig, - *, - subgoal: Subgoal, - node_scope: Optional[str], - query: str, - depth: int, - steps_out: Optional[List[Any]], - result: HarvestResult, - allowed_section_ids: Optional[Set[str]] = None, -) -> None: - from .nav_navigate import _apply_collect # late import avoids cycle - from .nav_token_budget import stamp_step_detail - - max_depth = max(0, int(getattr(config, "max_harvest_depth", 0) or 0)) - show_harvested = bool(config.is_checklist) - subgoal_dismissed = state.subgoal_dismissed_section_ids.get(subgoal.id, set()) - projection = build_map( - ts, - doc_id=state.doc_id, - query=query, - scope=node_scope, - config=config, - map_scores=state.map_scores, - collected_section_ids=state.collected_section_ids, - dismissed_section_ids=state.dismissed_section_ids | subgoal_dismissed, - highlight_ids=state.highlight_ids, - harvested_section_ids=state.harvested_owner_subgoal if show_harvested else None, - allowed_section_ids=allowed_section_ids, - ) - actions = build_legal_actions( - state, projection, config=config, depth=depth, ts=ts - ) - actionable = [a for a in actions if a.kind != ActionKind.FINISH] - if not actionable: - result.reason = result.reason or "no_legal_actions" - return - - collect_actions, dispatch_actions, confidence, reason, meta = harvest_policy_call( - ts, - state, - config, - subgoal=subgoal, - query=query, - projection=projection, - actions=actions, - depth=depth, - ) - result.n_policy_calls += 1 - if reason: - tag = node_scope or "root" - result.reason = f"{result.reason} | {tag}: {reason}" if result.reason else f"{tag}: {reason}" - - search_assets = list(meta.get("search_assets") or []) - result.search_assets.extend(search_assets) - - # Only the current node and its direct children are an actual decision at - # this call — deeper rows are shown for context inside the same - # full-depth map but belong to whichever recursive call lands on their - # own direct parent. Dismissing them here would pre-empt that call (e.g. - # a leaf two levels down would never get its own chance once its - # grandparent is DISPATCHed instead of collected). - # search_assets does not count as selecting a map node (F5). - depth_from_scope = {v.section_id: v.depth_from_scope for v in projection.tree_sections} - decided_now = [a for a in actionable if depth_from_scope.get(a.section_id or "", 99) <= 1] - selected_ids = {a.section_id for a in (*collect_actions, *dispatch_actions) if a.section_id} - unselected_ids = { - a.section_id for a in decided_now if a.section_id and a.section_id not in selected_ids - } - if unselected_ids: - state.subgoal_dismissed_section_ids.setdefault(subgoal.id, set()).update(unselected_ids) - - if steps_out is not None: - from ._compat import AgentStep # type: ignore - - observation = format_actionable_map_observation( - projection, list(actions), inline_summary=projection.scope is not None - ) - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="harvest", - detail=stamp_step_detail({ - "subgoal_id": subgoal.id, - "scope": node_scope, - "depth": depth, - "collect_ids": [a.action_id for a in collect_actions], - "dispatch_ids": [a.action_id for a in dispatch_actions], - "search_assets": search_assets, - "confidence": confidence, - "reason": reason, - "projection_chars": len(observation), - "legal_actions_preview": [a.prompt_line() for a in actions[:16]], - "visible_section_ids": [ - str(v.section_id) - for v in projection.tree_sections - if getattr(v, "section_id", None) - ], - **{k: v for k, v in meta.items() if k != "search_assets"}, - }), - ) - ) - - if collect_actions: - conf_by_sid = { - str(a.section_id): float(confidence.get(a.action_id.upper(), 0.0)) - for a in collect_actions - if a.section_id - } - primary = collect_actions[0] - primary.metadata = dict(primary.metadata or {}) - primary.metadata["batch_actions"] = collect_actions - primary.metadata["confidence_by_section"] = conf_by_sid - cdetail = _apply_collect(ts, state, primary, config) - new_roots = list(cdetail.get("collect_section_ids") or []) - result.new_section_ids.extend(new_roots) - if show_harvested: - for sid in new_roots: - state.harvested_owner_subgoal[sid] = subgoal.id - - if search_assets: - from .nav_assets import apply_search_assets - - n_added, asset_trace = apply_search_assets( - ts, - state, - config, - requests=search_assets, - default_scope=node_scope, - ) - result.n_assets_added += n_added - if steps_out is not None and asset_trace: - from ._compat import AgentStep # type: ignore - - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="search_assets", - detail=stamp_step_detail({ - "subgoal_id": subgoal.id, - "scope": node_scope, - "n_added": n_added, - "requests": asset_trace, - }), - ) - ) - - if dispatch_actions and depth >= max_depth: - result.max_depth_hit = True - return - - from .nav_address import next_dispatch_depth - - for act in dispatch_actions: - sid = str(act.section_id or "").strip() - if not sid or sid == str(node_scope or ""): - continue - result.visited_section_ids.append(sid) - before_collected = len(result.new_section_ids) - child_depth = next_dispatch_depth( - ts, - parent_doc_id=str(state.doc_id or ""), - parent_scope=str(node_scope or ""), - child_id=sid, - depth=depth, - ) - _harvest_node( - ts, - state, - config, - subgoal=subgoal, - node_scope=sid, - query=query, - depth=child_depth, - steps_out=steps_out, - result=result, - allowed_section_ids=allowed_section_ids, - ) - # Whole dispatched subtree came up empty: a dead end, not just "not - # yet explored" — dismiss it too so a later widen doesn't re-dispatch - # into the same branch. - if len(result.new_section_ids) == before_collected: - state.subgoal_dismissed_section_ids.setdefault(subgoal.id, set()).add(sid) diff --git a/deprecated/mapnav/nav/nav_hierarchy.py b/deprecated/mapnav/nav/nav_hierarchy.py deleted file mode 100644 index cabef82e0..000000000 --- a/deprecated/mapnav/nav/nav_hierarchy.py +++ /dev/null @@ -1,111 +0,0 @@ -"""In-memory hierarchy fixtures for map-nav tests. - -``ProviderToolSpace`` / ``NodeMeta`` / ``HierarchyProvider`` live in -``shared.services.retrieval.scoring.hierarchy``. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Dict, List, Optional, Sequence, Set, Tuple - -from shared.services.retrieval.scoring.hierarchy import NodeMeta - -@dataclass -class InMemoryNode: - section_id: str - title: str - content: str = "" - children: List[str] = field(default_factory=list) - - -class InMemoryHierarchyProvider: - """Minimal reference ``HierarchyProvider``: no scoring, no ToolSpace. - - Built directly from a ``{doc_id: [InMemoryNode, ...]}`` map plus a - ``{doc_id: [root_section_id, ...]}`` map — the "hierarchy + summary is - enough" claim's simplest possible witness. - """ - - def __init__( - self, - *, - roots_by_doc: Dict[str, Sequence[str]], - nodes: Dict[str, InMemoryNode], - summaries: Optional[Dict[str, str]] = None, - ) -> None: - self._roots_by_doc = {k: list(v) for k, v in roots_by_doc.items()} - self._nodes = dict(nodes) - self._summaries = dict(summaries or {}) - self._parent: Dict[str, str] = {} - for node in self._nodes.values(): - for child_id in node.children: - self._parent[child_id] = node.section_id - self._owner: Dict[str, str] = {} - for doc_id, root_ids in self._roots_by_doc.items(): - stack = list(root_ids) - while stack: - sid = stack.pop() - if sid in self._owner: - continue - self._owner[sid] = doc_id - node = self._nodes.get(sid) - if node: - stack.extend(node.children) - - def owner_document(self, node_id: str) -> Optional[str]: - return self._owner.get(str(node_id or "").strip()) - - def roots(self, doc_id: str) -> Sequence[str]: - return list(self._roots_by_doc.get(doc_id, ())) - - def children(self, section_id: str) -> Sequence[str]: - node = self._nodes.get(section_id) - return list(node.children) if node else [] - - def node_meta(self, section_id: str) -> NodeMeta: - node = self._nodes.get(section_id) - if node is None: - return NodeMeta() - return NodeMeta( - title=node.title, - summary=self._summaries.get(section_id, ""), - has_children=bool(node.children), - ) - - def parent_id(self, section_id: str) -> Optional[str]: - return self._parent.get(section_id) - - def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: - ancestors: Set[str] = set() - cur = self._parent.get(section_id) - while cur: - ancestors.add(cur) - cur = self._parent.get(cur) - descendants: Set[str] = set() - stack = list(self.children(section_id)) - while stack: - cid = stack.pop() - if cid in descendants: - continue - descendants.add(cid) - stack.extend(self.children(cid)) - return ancestors, descendants - - def content(self, section_id: str) -> str: - node = self._nodes.get(section_id) - if node is None: - return "" - parts: List[str] = [] - - def walk(sid: str) -> None: - cur = self._nodes.get(sid) - if cur is None: - return - if cur.content: - parts.append(cur.content) - for cid in cur.children: - walk(cid) - - walk(section_id) - return "\n".join(parts) diff --git a/deprecated/mapnav/nav/nav_knowhere.py b/deprecated/mapnav/nav/nav_knowhere.py deleted file mode 100644 index 90d5cf3d9..000000000 --- a/deprecated/mapnav/nav/nav_knowhere.py +++ /dev/null @@ -1,983 +0,0 @@ -"""Knowhere-native hierarchy provider leftovers for MAP-NAV. - -Publication types and the eager ``KnowhereProvider`` live in -``shared.services.retrieval.scoring``. This module keeps the episode-local -lazy/DB loaders used by map-nav snapshot loading. -""" - -from __future__ import annotations - -import os -import logging -import time -from hashlib import sha256 -from typing import ( - Any, - Dict, - Iterable, - List, - Mapping, - Optional, - Protocol, - Sequence, - Set, - Tuple, -) - -from shared.services.retrieval.nav.nav_address import NavLevel -from shared.services.retrieval.scoring.hierarchy import NodeMeta -from shared.services.retrieval.scoring.knowhere_hybrid import ( - MAP_UNIT_INDEX_FORMAT_VERSION, - PersistedScoreCorpus, - PersistedScoreUnit, - tokenize_query_for_ranker, -) -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, - UnitRow, - is_root_section_path, -) - -_DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" -_MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") -_logger = logging.getLogger(__name__) - -def is_root_section(provider_or_ts: Any, section_id: str) -> bool: - """True when ``section_id`` is a Root container (raw path, not normalized).""" - sid = str(section_id or "").strip() - if not sid: - return False - path_fn = getattr(provider_or_ts, "section_path", None) - if callable(path_fn): - return is_root_section_path(path_fn(sid)) - provider = getattr(provider_or_ts, "_provider", None) - path_fn = getattr(provider, "section_path", None) if provider is not None else None - if callable(path_fn): - return is_root_section_path(path_fn(sid)) - return False - - -def knowhere_database_url() -> str: - configured = ( - str(os.environ.get("KNOWHERE_DATABASE_URL") or "").strip() - or str(os.environ.get("DATABASE_URL") or "").strip() - ) - if configured: - # ``ReadOnlyChunkStore`` uses psycopg2's native connector, which - # accepts libpq URLs but not SQLAlchemy's ``+driver`` suffix. - return configured.replace("postgresql+asyncpg://", "postgresql://", 1).replace( - "postgresql+psycopg2://", "postgresql://", 1 - ) - return _DEFAULT_DSN - - -class ChunkStore(Protocol): - def load_chunk_reference_metadata( - self, - document_id: str, - chunk_id: str, - ) -> Optional[Mapping[str, Any]]: - raise NotImplementedError - - def load_persisted_score_corpus( - self, - document_ids: Sequence[str], - allowed_section_ids_by_document: Mapping[str, Sequence[str]], - queries: Sequence[str], - ) -> Optional[PersistedScoreCorpus]: - raise NotImplementedError - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> List[UnitRow]: - raise NotImplementedError - - def close(self) -> None: - raise NotImplementedError - - -class ReadOnlyChunkStore: - """Episode-local, revision-pinned loader for lazy map-nav chunks.""" - - def __init__( - self, - *, - dsn: str, - revisions: Dict[str, str], - excluded_sections: Optional[Iterable[Tuple[str, str]]] = None, - ) -> None: - self._dsn = str(dsn) - self._revisions = dict(revisions) - self._excluded_sections = set(excluded_sections or ()) - self._conn: Optional[_SyncConnection] = None - self._score_unit_rows_cache: dict[ - tuple[tuple[str, str], ...], list[dict[str, object]] - ] = {} - self._score_average_idf_cache: dict[ - tuple[tuple[str, str], ...], tuple[float, float] - ] = {} - - def _connection(self) -> "_SyncConnection": - if self._conn is None: - self._conn = _connect(self._dsn) - self._conn.set_session(readonly=True, autocommit=True) - return self._conn - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> List[UnitRow]: - doc_id = str(document_id).strip() - sid = str(section_id).strip() - job_result_id = self._revisions.get(doc_id) - if ( - not doc_id - or not sid - or not job_result_id - or (doc_id, sid) in self._excluded_sections - ): - return [] - cur = self._connection().cursor() - try: - ids = [ - str(chunk_id).strip() - for chunk_id in extra_chunk_ids - if str(chunk_id).strip() - ] - if ids: - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s " - "AND (section_id = %s OR chunk_id = ANY(%s)) " - "ORDER BY sort_order, chunk_id, id", - (doc_id, job_result_id, sid, ids), - ) - else: - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s AND section_id = %s " - "ORDER BY sort_order, chunk_id, id", - (doc_id, job_result_id, sid), - ) - return [_unit_from_row(row) for row in cur.fetchall()] - finally: - cur.close() - - def load_chunk_reference_metadata( - self, - document_id: str, - chunk_id: str, - ) -> Optional[Mapping[str, Any]]: - """Resolve deferred reference fields for one selected chunk.""" - doc_id = str(document_id).strip() - cid = str(chunk_id).strip() - job_result_id = self._revisions.get(doc_id) - if not doc_id or not cid or not job_result_id: - return None - cur = self._connection().cursor() - try: - cur.execute( - "SELECT file_path FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s AND chunk_id = %s " - "ORDER BY sort_order DESC, id DESC LIMIT 1", - (doc_id, job_result_id, cid), - ) - row = cur.fetchone() - return {"file_path": str(row[0] or "") or None} if row else None - finally: - cur.close() - - def load_persisted_score_corpus( - self, - document_ids: Sequence[str], - allowed_section_ids_by_document: Mapping[str, Sequence[str]], - queries: Sequence[str], - ) -> Optional[PersistedScoreCorpus]: - """Load query-token score inputs from map-unit tables when indexed. - - Units and lengths come from ``document_map_units``. Frequencies come from - ``document_map_unit_tokens`` filtered to the query tokens. Average IDF - comes from ``document_map_unit_indexes`` (written at index time). - """ - from shared.services.retrieval.scoring.persisted_score_load import ( - build_channel_bm25_stats, - combine_average_idf, - ) - - revisions = [ - (document_id, self._revisions[document_id]) - for raw_document_id in document_ids - if (document_id := str(raw_document_id).strip()) in self._revisions - ] - if not revisions or len(revisions) != len(document_ids): - return None - values_sql = ", ".join(["(%s, %s)"] * len(revisions)) - revision_params: List[object] = [ - value for revision in revisions for value in revision - ] - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - query_tokens = list( - dict.fromkeys( - token - for query in unique_queries - for token in tokenize_query_for_ranker(query) - ) - ) - query_token_hashes = [ - sha256(token.encode("utf-8")).hexdigest() for token in query_tokens - ] - cur = self._connection().cursor() - try: - revision_key = tuple(revisions) - stage_started = time.perf_counter() - try: - cur.execute( - "SELECT indexes.document_id, indexes.job_result_id, " - "indexes.format_version, indexes.unit_count, " - "indexes.average_idf_path, indexes.average_idf_content, " - "indexes.path_document_count, indexes.path_total_length, " - "indexes.content_document_count, indexes.content_total_length " - "FROM document_map_unit_indexes AS indexes " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON indexes.document_id = revisions.document_id " - "AND indexes.job_result_id = revisions.job_result_id", - revision_params, - ) - except Exception as exc: - if getattr(exc, "pgcode", None) in {"42P01", "42703"}: - return None - raise - index_rows = list(cur.fetchall()) - _logger.info( - "retrieval map-index load stage=indexes seconds=%.3f rows=%d", - time.perf_counter() - stage_started, - len(index_rows), - ) - if len(index_rows) != len(revisions) or any( - len(row) < 10 or int(row[2]) != MAP_UNIT_INDEX_FORMAT_VERSION - or any(value is None for value in row[6:10]) - for row in index_rows - ): - return None - - if revision_key in self._score_average_idf_cache: - average_idf_path, average_idf_content = self._score_average_idf_cache[ - revision_key - ] - else: - average_idf_path = combine_average_idf( - [(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows] - ) - average_idf_content = combine_average_idf( - [ - (float(row[5] or 0.0), int(row[3] or 0)) - for row in index_rows - ] - ) - self._score_average_idf_cache[revision_key] = ( - average_idf_path, - average_idf_content, - ) - - allowed_by_document = { - str(document_id): {str(section_id) for section_id in section_ids} - for document_id, section_ids in allowed_section_ids_by_document.items() - } - allowed_pairs_set = { - (document_id, section_id) - for document_id, section_ids in allowed_by_document.items() - for section_id in section_ids - } - cur.execute( - "SELECT sections.document_id, sections.job_result_id, count(*) " - "FROM document_sections AS sections " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON sections.document_id = revisions.document_id " - "AND sections.job_result_id = revisions.job_result_id " - "GROUP BY sections.document_id, sections.job_result_id", - revision_params, - ) - section_counts = { - (str(document_id), str(job_result_id)): int(count) - for document_id, job_result_id, count in cur.fetchall() - } - has_complete_section_scope = all( - len(allowed_by_document.get(document_id, set())) - == section_counts.get((document_id, job_result_id), 0) - for document_id, job_result_id in revisions - ) - all_unit_rows = self._score_unit_rows_cache.get(revision_key) - if has_complete_section_scope and query_token_hashes: - stage_started = time.perf_counter() - cur.execute( - "WITH matching_tokens AS MATERIALIZED (" - "SELECT DISTINCT map_unit_id FROM document_map_unit_tokens " - "WHERE channel = ANY(%s) AND token_hash = ANY(%s)" - "), scoped_units AS MATERIALIZED (" - f"SELECT units.id, units.document_id, units.unit_id, units.section_id, " - "units.path_token_count, units.content_token_count " - "FROM document_map_units AS units " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id" - ") SELECT scoped_units.id, scoped_units.document_id, " - "scoped_units.unit_id, scoped_units.section_id, " - "scoped_units.path_token_count, scoped_units.content_token_count " - "FROM matching_tokens JOIN scoped_units " - "ON scoped_units.id = matching_tokens.map_unit_id", - [ - list(_MAP_SCORE_CHANNELS), - list(query_token_hashes), - *revision_params, - ], - ) - unit_rows = [ - { - "map_unit_id": str(row[0]), - "document_id": str(row[1]), - "unit_id": str(row[2]), - "section_id": str(row[3] or ""), - "path_token_count": int(row[4] or 0), - "content_token_count": int(row[5] or 0), - } - for row in cur.fetchall() - ] - _logger.info( - "retrieval map-index load stage=units-selective seconds=%.3f rows=%d", - time.perf_counter() - stage_started, - len(unit_rows), - ) - all_unit_rows = None - elif all_unit_rows is None: - stage_started = time.perf_counter() - cur.execute( - "SELECT units.id, units.document_id, units.unit_id, " - "units.section_id, units.path_token_count, units.content_token_count " - "FROM document_map_units AS units " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id " - "ORDER BY units.document_id, units.sort_order, units.unit_id", - revision_params, - ) - all_unit_rows = [ - { - "map_unit_id": str(row[0]), - "document_id": str(row[1]), - "unit_id": str(row[2]), - "section_id": str(row[3] or ""), - "path_token_count": int(row[4] or 0), - "content_token_count": int(row[5] or 0), - } - for row in cur.fetchall() - ] - expected_unit_count = sum(int(row[3] or 0) for row in index_rows) - if expected_unit_count and len(all_unit_rows) != expected_unit_count: - return None - self._score_unit_rows_cache[revision_key] = all_unit_rows - _logger.info( - "retrieval map-index load stage=units seconds=%.3f rows=%d cache_hit=%s", - time.perf_counter() - stage_started, - len(all_unit_rows), - False, - ) - else: - _logger.info( - "retrieval map-index load stage=units seconds=0.000 rows=%d cache_hit=%s", - len(all_unit_rows), - True, - ) - - if not has_complete_section_scope or not query_token_hashes: - unit_rows = [ - row - for row in all_unit_rows or [] - if (str(row["document_id"]), str(row["section_id"])) - in allowed_pairs_set - ] - frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} - if unit_rows and query_tokens: - stage_started = time.perf_counter() - # Keep the episode scope explicit without materializing every - # token row matching a common query term. PostgreSQL can choose - # either the scoped-unit side or the channel/token_hash-leading - # index, while the scope still limits results to the pinned - # revision and allowed sections. - allowed_map_unit_ids = [str(row["map_unit_id"]) for row in unit_rows] - cur.execute( - "WITH scoped_units AS MATERIALIZED (" - "SELECT unnest(%s::text[]) AS map_unit_id" - ") " - "SELECT tokens.map_unit_id, tokens.channel, tokens.token, " - "tokens.frequency " - "FROM document_map_unit_tokens AS tokens " - "JOIN scoped_units " - "ON scoped_units.map_unit_id = tokens.map_unit_id " - "WHERE tokens.channel = ANY(%s) " - "AND tokens.token_hash = ANY(%s)", - [ - allowed_map_unit_ids, - list(_MAP_SCORE_CHANNELS), - list(query_token_hashes), - ], - ) - for map_unit_id, channel, token, frequency in cur.fetchall(): - frequencies.setdefault((str(map_unit_id), str(channel)), {})[ - str(token) - ] = int(frequency) - _logger.info( - "retrieval map-index load stage=frequencies seconds=%.3f units=%d tokens=%d", - time.perf_counter() - stage_started, - len(unit_rows), - len(query_tokens), - ) - - _logger.info( - "retrieval map-index load stage=complete units=%d queries=%d", - len(unit_rows), - len(unique_queries), - ) - path_stats = build_channel_bm25_stats( - unit_rows=unit_rows, - map_unit_id_field="map_unit_id", - length_field="path_token_count", - channel="path", - query_tokens=query_tokens, - frequencies=frequencies, - average_idf=average_idf_path, - document_count_override=( - sum(int(row[6] or 0) for row in index_rows) - if has_complete_section_scope and query_token_hashes - else None - ), - total_length_override=( - sum(int(row[7] or 0) for row in index_rows) - if has_complete_section_scope and query_token_hashes - else None - ), - ) - content_stats = build_channel_bm25_stats( - unit_rows=unit_rows, - map_unit_id_field="map_unit_id", - length_field="content_token_count", - channel="content", - query_tokens=query_tokens, - frequencies=frequencies, - average_idf=average_idf_content, - document_count_override=( - sum(int(row[8] or 0) for row in index_rows) - if has_complete_section_scope and query_token_hashes - else None - ), - total_length_override=( - sum(int(row[9] or 0) for row in index_rows) - if has_complete_section_scope and query_token_hashes - else None - ), - ) - return PersistedScoreCorpus( - units=[ - PersistedScoreUnit( - unit_id=str(row["unit_id"]), - path_length=int(row["path_token_count"]), - content_length=int(row["content_token_count"]), - path_frequencies=frequencies.get( - (str(row["map_unit_id"]), "path"), {} - ), - content_frequencies=frequencies.get( - (str(row["map_unit_id"]), "content"), {} - ), - ) - for row in unit_rows - # Only units with at least one query-token frequency can - # score > 0; zero-frequency units are implicit 0 and are - # never materialized for BM25. - if frequencies.get((str(row["map_unit_id"]), "path")) - or frequencies.get((str(row["map_unit_id"]), "content")) - ], - path_stats=path_stats, - content_stats=content_stats, - ) - finally: - cur.close() - - def close(self) -> None: - if self._conn is not None: - self._conn.close() - self._conn = None - - -class _SyncCursor(Protocol): - def execute(self, query: str, params: Sequence[object]) -> None: - raise NotImplementedError - - def fetchall(self) -> Sequence[Sequence[object]]: - raise NotImplementedError - - def fetchone(self) -> Optional[Sequence[object]]: - raise NotImplementedError - - def close(self) -> None: - raise NotImplementedError - - -class _SyncConnection(Protocol): - def set_session(self, *, readonly: bool, autocommit: bool) -> None: - raise NotImplementedError - - def cursor(self) -> _SyncCursor: - raise NotImplementedError - - def close(self) -> None: - raise NotImplementedError - - -def _unit_from_row(row: Sequence[object]) -> UnitRow: - return UnitRow( - chunk_id=str(row[0] or ""), - section_id=str(row[1]) if row[1] else None, - 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]), - ) - - -class LazyKnowhereProvider(KnowhereProvider): - """Hierarchy provider that loads full chunk rows only on first access.""" - - def __init__( - self, - *, - doc_id: str, - sections: Sequence[SectionRow], - chunk_store: ChunkStore, - known_chunk_ids: Sequence[str], - root_asset_ids: Sequence[str] = (), - remounted_assets_by_section: Optional[Dict[str, Sequence[str]]] = None, - ) -> None: - super().__init__( - doc_id=doc_id, - sections=sections, - units=(), - lazy_loader=lambda section_id: chunk_store.load_section_units( - doc_id, - section_id, - (remounted_assets_by_section or {}).get(section_id, ()), - ), - known_chunk_ids=known_chunk_ids, - ) - self._chunk_store = chunk_store - self._root_asset_ids = {str(chunk_id) for chunk_id in root_asset_ids} - - def _ensure_section_loaded(self, section_id: str) -> None: - super()._ensure_section_loaded(section_id) - if ( - section_id in self._units_by_section - and self._root_asset_ids - and is_root_section_path(self.section_path(section_id)) - ): - self._units_by_section[section_id] = [ - unit - for unit in self._units_by_section[section_id] - if unit.chunk_id not in self._root_asset_ids - ] - - def close(self) -> None: - self._chunk_store.close() - - -def _connect(dsn: str) -> _SyncConnection: - import psycopg2 - - return psycopg2.connect(dsn) - - -def _as_meta(raw: Any) -> Dict[str, Any]: - if isinstance(raw, dict): - return dict(raw) - if isinstance(raw, str) and raw.strip(): - import json - - try: - obj = json.loads(raw) - except Exception: - return {} - return dict(obj) if isinstance(obj, dict) else {} - return {} - - -def load_document_from_db( - document_id: str, - *, - dsn: Optional[str] = None, -) -> KnowhereProvider: - """Load one document's current revision into a ``KnowhereProvider``.""" - doc_id = str(document_id or "").strip() - if not doc_id: - raise ValueError("document_id is required") - url = dsn or knowhere_database_url() - conn = _connect(url) - try: - with conn.cursor() as cur: - cur.execute( - "SELECT document_id, current_job_result_id, source_file_name " - "FROM documents WHERE document_id = %s AND status = 'active'", - (doc_id,), - ) - row = cur.fetchone() - if row is None: - raise ValueError(f"active document not found: {doc_id}") - job_result_id = row[1] - if not job_result_id: - raise ValueError(f"document {doc_id} has no current_job_result_id") - - cur.execute( - "SELECT section_id, parent_section_id, section_path, section_title, " - "section_level, summary, sort_order " - "FROM document_sections " - "WHERE document_id = %s AND job_result_id = %s " - "ORDER BY sort_order, section_id", - (doc_id, job_result_id), - ) - sections = [ - SectionRow( - section_id=str(r[0]), - parent_section_id=str(r[1]) if r[1] else None, - section_path=str(r[2] or ""), - section_title=str(r[3] or "").strip(), - section_level=int(r[4] or 0), - summary=str(r[5] or "").strip(), - sort_order=int(r[6] or 0), - ) - for r in cur.fetchall() - ] - - cur.execute( - "SELECT chunk_id, section_id, chunk_type, content, sort_order, " - "source_chunk_path, file_path, chunk_metadata " - "FROM document_chunks " - "WHERE document_id = %s AND job_result_id = %s " - "ORDER BY sort_order, chunk_id", - (doc_id, job_result_id), - ) - units = [ - UnitRow( - chunk_id=str(r[0] or ""), - section_id=str(r[1]) if r[1] else None, - chunk_type=str(r[2] or "text"), - content=str(r[3] or ""), - sort_order=int(r[4] or 0), - source_chunk_path=str(r[5] or ""), - file_path=str(r[6] or ""), - metadata=_as_meta(r[7]), - ) - for r in cur.fetchall() - ] - finally: - conn.close() - return KnowhereProvider(doc_id=doc_id, sections=sections, units=units) - - -def load_namespace_from_db( - *, - namespace: str, - document_ids: Optional[Sequence[str]] = None, - titles: Optional[Dict[str, str]] = None, - dsn: Optional[str] = None, -) -> NamespaceKnowhereProvider: - """Load active documents in a namespace (or an explicit id list).""" - ns = str(namespace or "").strip() - url = dsn or knowhere_database_url() - wanted = [str(d).strip() for d in (document_ids or ()) if str(d).strip()] - if not wanted: - if not ns: - raise ValueError("namespace or document_ids is required") - conn = _connect(url) - try: - with conn.cursor() as cur: - cur.execute( - "SELECT document_id, source_file_name FROM documents " - "WHERE namespace = %s AND status = 'active' " - "ORDER BY document_id", - (ns,), - ) - rows = cur.fetchall() - wanted = [str(r[0]) for r in rows] - auto_titles = { - str(r[0]): str(r[1] or r[0]).strip() for r in rows if r[0] - } - finally: - conn.close() - else: - auto_titles = {} - - if not wanted: - raise ValueError(f"no active documents for namespace={ns!r}") - - providers = [load_document_from_db(did, dsn=url) for did in wanted] - merged_titles = dict(auto_titles) - if titles: - merged_titles.update( - {str(k): str(v) for k, v in titles.items() if str(k).strip()} - ) - return NamespaceKnowhereProvider(providers, titles=merged_titles or None) - - -class NamespaceKnowhereProvider: - """Multi-document provider: document ids are DISPATCH-only map nodes. - - Namespace root is empty scope (not a node). ``roots("")`` returns the - document ids; ``children(document_id)`` returns that document's section - roots. Section/chunk identity and ownership stay on the real keys. - """ - - def __init__( - self, - providers: Sequence[KnowhereProvider], - *, - titles: Optional[Dict[str, str]] = None, - chunk_owner_by_id: Optional[Dict[str, str]] = None, - ) -> None: - self._docs: Dict[str, KnowhereProvider] = { - p.doc_id: p for p in providers if p.doc_id - } - if not self._docs: - raise ValueError("NamespaceKnowhereProvider requires at least one document") - self._titles = { - did: str((titles or {}).get(did) or did).strip() or did - for did in self._docs - } - self._section_owner: Dict[str, str] = {} - self._chunk_owner: Dict[str, str] = {} - for doc_id, provider in self._docs.items(): - for sid in provider.all_section_ids(): - self._section_owner[sid] = doc_id - if chunk_owner_by_id: - self._chunk_owner.update( - { - str(chunk_id): str(doc_id) - for chunk_id, doc_id in chunk_owner_by_id.items() - if str(chunk_id).strip() and str(doc_id).strip() - } - ) - else: - for doc_id, provider in self._docs.items(): - for sid in provider.all_section_ids(): - for unit in provider.self_units(sid): - if unit.chunk_id: - self._chunk_owner[unit.chunk_id] = doc_id - - def document_ids(self) -> List[str]: - return list(self._docs) - - def close(self) -> None: - for provider in self._docs.values(): - close = getattr(provider, "close", None) - if callable(close): - close() - - def load_persisted_score_corpus( - self, - doc_ids: Sequence[str], - queries: Sequence[str], - ) -> Optional[PersistedScoreCorpus]: - """Return the index projection only when all documents share one store.""" - providers = [ - self._docs[document_id] - for raw_document_id in doc_ids - if (document_id := str(raw_document_id).strip()) in self._docs - ] - if len(providers) != len(doc_ids) or not all( - isinstance(provider, LazyKnowhereProvider) for provider in providers - ): - return None - lazy_providers = [ - provider - for provider in providers - if isinstance(provider, LazyKnowhereProvider) - ] - stores = { - id(provider._chunk_store): provider._chunk_store - for provider in lazy_providers - } - if len(stores) != 1: - return None - store = next(iter(stores.values())) - loader = getattr(store, "load_persisted_score_corpus", None) - if not callable(loader): - return None - return loader( - [provider.doc_id for provider in lazy_providers], - {provider.doc_id: list(provider._sections) for provider in lazy_providers}, - queries, - ) - - def address_level(self, node_id: str) -> Optional[NavLevel]: - sid = str(node_id or "").strip() - if not sid: - return NavLevel.NAMESPACE - if sid in self._docs: - return NavLevel.DOCUMENT - if sid in self._section_owner: - return NavLevel.SECTION - if sid in self._chunk_owner: - return NavLevel.CHUNK - return None - - def owner_document(self, node_id: str) -> Optional[str]: - sid = str(node_id or "").strip() - if not sid: - return None - if sid in self._docs: - return sid - return self._section_owner.get(sid) or self._chunk_owner.get(sid) - - def roots(self, doc_id: str) -> Sequence[str]: - key = str(doc_id or "").strip() - if not key: - return list(self._docs) - provider = self._docs.get(key) - return list(provider.roots(key)) if provider else [] - - def children(self, section_id: str) -> Sequence[str]: - sid = str(section_id or "").strip() - if sid in self._docs: - return list(self._docs[sid].roots(sid)) - owner = self._section_owner.get(sid) - if not owner: - return [] - return list(self._docs[owner].children(sid)) - - def node_meta(self, section_id: str) -> NodeMeta: - sid = str(section_id or "").strip() - if sid in self._docs: - provider = self._docs[sid] - return NodeMeta( - title=self._titles.get(sid, sid), - summary="", - has_children=bool(provider.roots(sid)), - ) - owner = self._section_owner.get(sid) - if not owner: - return NodeMeta() - return self._docs[owner].node_meta(sid) - - def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: - sid = str(section_id or "").strip() - if sid in self._docs: - descendants = set(self._docs[sid].all_section_ids()) - return set(), descendants - owner = self._section_owner.get(sid) - if not owner: - return set(), set() - ancestors, descendants = self._docs[owner].relations(sid) - row_parent = self._docs[owner].parent_id(sid) - if row_parent is None: - ancestors = set(ancestors) | {owner} - return ancestors, descendants - - def content(self, section_id: str) -> str: - sid = str(section_id or "").strip() - if sid in self._docs: - provider = self._docs[sid] - return "\n".join( - provider.content(root) - for root in provider.roots(sid) - if provider.content(root) - ) - owner = self._section_owner.get(sid) - if not owner: - return "" - return self._docs[owner].content(sid) - - def self_units(self, section_id: str) -> List[UnitRow]: - sid = str(section_id or "").strip() - if sid in self._docs: - return [] - owner = self._section_owner.get(sid) - if not owner: - return [] - return self._docs[owner].self_units(sid) - - def leaf_ids(self, section_id: str) -> List[str]: - sid = str(section_id or "").strip() - if sid in self._docs: - out: List[str] = [] - for root in self._docs[sid].roots(sid): - out.extend(self._docs[sid].leaf_ids(root)) - return out - owner = self._section_owner.get(sid) - if not owner: - return [] - return self._docs[owner].leaf_ids(sid) - - def path_titles(self, section_id: str) -> str: - sid = str(section_id or "").strip() - if sid in self._docs: - return self._titles.get(sid, sid) - owner = self._section_owner.get(sid) - if not owner: - return "" - title = self._docs[owner].path_titles(sid) - doc_title = self._titles.get(owner, owner) - return f"{doc_title} / {title}" if title else doc_title - - def parent_id(self, section_id: str) -> Optional[str]: - sid = str(section_id or "").strip() - if sid in self._docs: - return None - owner = self._section_owner.get(sid) - if not owner: - return None - parent = self._docs[owner].parent_id(sid) - return parent if parent is not None else owner - - def section_path(self, section_id: str) -> str: - sid = str(section_id or "").strip() - owner = self._section_owner.get(sid) - if not owner: - return "" - return self._docs[owner].section_path(sid) - - def resolve_path(self, path: str, doc_id: str = "") -> Optional[str]: - did = str(doc_id or "").strip() - if did and did in self._docs: - return self._docs[did].resolve_path(path) - for provider in self._docs.values(): - hit = provider.resolve_path(path) - if hit: - return hit - return None - - def unit_text(self, unit: UnitRow) -> str: - owner = self._chunk_owner.get(unit.chunk_id) or self._section_owner.get( - str(unit.section_id or "") - ) - if not owner: - return "" - return self._docs[owner].unit_text(unit) - - def summaries(self) -> Dict[str, str]: - out: Dict[str, str] = {} - for provider in self._docs.values(): - out.update(provider.summaries()) - return out - - def all_section_ids(self) -> List[str]: - out: List[str] = [] - for provider in self._docs.values(): - out.extend(provider.all_section_ids()) - return out diff --git a/deprecated/mapnav/nav/nav_llm.py b/deprecated/mapnav/nav/nav_llm.py deleted file mode 100644 index 4cd0bfa81..000000000 --- a/deprecated/mapnav/nav/nav_llm.py +++ /dev/null @@ -1,266 +0,0 @@ -"""Boundary LLM adapter for MAP-NAV (Knowhere-aligned call shape). - -Nav modules call ``nav_chat`` / ``resolve_nav_model`` only — they must not read -``DS_KEY`` / ``OPENAI_API_KEY`` themselves. Credentials are resolved here via -``resolve_chat_credentials`` (deepseek-* → ``DS_KEY``/``DS_URL``, else -``OPENAI_*``), matching Knowhere's chat credential split (deepseek → DS_*, else OPENAI_*). - -Thinking policy (DeepSeek V4 defaults thinking ON if omitted): - -- ``action`` (harvest / refine / verify / score): - always disabled — short JSON under ``llm_max_tokens`` (often 256). -- ``planner`` (plan_query / replan only): episode-bound - ``NavConfig.planner_thinking``, else ``NAV_PLANNER_THINKING`` for EXP - scripts; unset → disabled. When enabled, callers should use - ``planner_output_max_tokens``. ``plan_control`` uses ``action`` (thinking off). - -Migration to Knowhere: inject the production callable with -``set_nav_chat_backend`` (wrap ``llm_fn``); leave nav call sites unchanged. -""" - -from __future__ import annotations - -import logging -import os -import time -from contextlib import contextmanager -from contextvars import ContextVar -from typing import Any, Callable, Dict, Iterator, Optional, Sequence - -# Same return shape as ``cached_chat_completion``: -# ``{content, reasoning_content?, usage, cache_hit?}``. -NavChatBackend = Callable[..., Dict[str, Any]] - -_backend: Optional[NavChatBackend] = None -_logger = logging.getLogger(__name__) - -_DS_DEFAULT_MODEL = "deepseek-v4-flash" -_DEFAULT_PLANNER_THINK_MAX = 16384 - -_runtime_planner_thinking: ContextVar[Optional[str]] = ContextVar( - "nav_runtime_planner_thinking", default=None -) -_runtime_planner_think_max: ContextVar[Optional[int]] = ContextVar( - "nav_runtime_planner_think_max", default=None -) - - -def set_nav_chat_backend(fn: Optional[NavChatBackend]) -> None: - """Install or clear an injected chat backend (tests / Knowhere llm_fn).""" - global _backend - _backend = fn - - -def get_nav_chat_backend() -> Optional[NavChatBackend]: - return _backend - - -@contextmanager -def nav_llm_runtime( - *, - planner_thinking: str = "", - planner_think_max_tokens: int = 0, -) -> Iterator[None]: - """Bind planner thinking knobs for one episode (config values, not process env).""" - think = str(planner_thinking or "").strip() or None - think_max = int(planner_think_max_tokens or 0) - t_tok = _runtime_planner_thinking.set(think) - m_tok = _runtime_planner_think_max.set(think_max if think_max > 0 else None) - try: - yield - finally: - _runtime_planner_think_max.reset(m_tok) - _runtime_planner_thinking.reset(t_tok) - - -def resolve_nav_model( - *, - model: str = "", - model_env: str = "", - fallback_envs: Sequence[str] = (), -) -> str: - """Prefer explicit model string; else env keys (EXP); else deepseek default.""" - preferred = str(model or "").strip() - if preferred: - return preferred - - from ._compat import load_llm_env # type: ignore - - load_llm_env() - for name in (model_env, *fallback_envs): - token = str(name or "").strip() - if not token: - continue - val = os.environ.get(token, "").strip() - if val: - return val - return _DS_DEFAULT_MODEL - - -def resolve_nav_thinking_mode(*, role: str = "action") -> str: - """Return ``enabled`` or ``disabled`` — never leave API auto (DS default ON).""" - from ._compat import ( # type: ignore - load_llm_env, - resolve_thinking_mode, - ) - - if role != "planner": - return "disabled" - - bound = _runtime_planner_thinking.get() - if bound in ("enabled", "disabled"): - return bound - - load_llm_env() - mode = resolve_thinking_mode( - os.environ.get("NAV_PLANNER_THINKING", "").strip() or None - ) - return mode if mode in ("enabled", "disabled") else "disabled" - - -def nav_thinking_extra(*, role: str = "action", model: str = "") -> Dict[str, Any]: - """Build cache-aware thinking ``extra`` for ``nav_chat`` / ``cached_chat_completion``.""" - from ._compat import chat_thinking_extra # type: ignore - - return chat_thinking_extra( - mode=resolve_nav_thinking_mode(role=role), - model=model, - ) - - -def planner_output_max_tokens(base: int) -> int: - """Raise planner/control max_tokens when planner thinking is enabled.""" - max_tokens = max(256, int(base or 0)) - if resolve_nav_thinking_mode(role="planner") != "enabled": - return max_tokens - bound = _runtime_planner_think_max.get() - if bound is not None and int(bound) > 0: - think_max = int(bound) - else: - think_max = int(os.environ.get("NAV_PLANNER_THINK_MAX_TOKENS", "").strip() or "0") - if think_max <= 0: - think_max = _DEFAULT_PLANNER_THINK_MAX - return max(max_tokens, think_max) - - -def _merge_thinking_extra( - extra: Optional[Dict[str, Any]], - *, - role: str, - model: str, -) -> Dict[str, Any]: - think = nav_thinking_extra(role=role, model=model) - merged = dict(extra or {}) - body = dict(merged.get("extra_body") or {}) - body.update(think.get("extra_body") or {}) - merged.update(think) - if body: - merged["extra_body"] = body - return merged - - -def nav_chat( - *, - purpose: str, - messages: Sequence[Dict[str, str]], - model: str, - temperature: float = 0.0, - max_tokens: Optional[int] = None, - response_format: Optional[Dict[str, Any]] = None, - extra: Optional[Dict[str, Any]] = None, - thinking_role: str = "action", - context: str = "Nav", - api_key_env: str = "", - base_url_env: str = "", - timeout: float = 60.0, - usage_tag: str = "", -) -> Dict[str, Any]: - """Single entry for nav LLM calls. Credentials stay inside this boundary. - - Token hard-stop: refuse the call when the episode/process budget is already - exhausted (raises ``NavTokenLimit``). Successful calls credit usage to the - active episode counter when ``nav_token_episode()`` is entered. - """ - from .nav_token_budget import ( - NavTokenLimit, - nav_token_budget_exhausted, - nav_token_limit, - nav_tokens_used, - record_episode_tokens, - ) - - if nav_token_budget_exhausted(): - raise NavTokenLimit(used=nav_tokens_used(), limit=nav_token_limit()) - - merged_extra = _merge_thinking_extra(extra, role=thinking_role, model=model) - call_started = time.perf_counter() - - if _backend is not None: - try: - result = _backend( - purpose=purpose, - messages=list(messages), - model=model, - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - thinking_role=thinking_role, - context=context, - api_key_env=api_key_env, - base_url_env=base_url_env, - timeout=timeout, - usage_tag=usage_tag, - ) - record_episode_tokens((result or {}).get("usage")) - return result - finally: - _logger.info( - "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", - purpose, - thinking_role, - time.perf_counter() - call_started, - ) - - from ._compat import cached_chat_completion # type: ignore - from ._compat import ( # type: ignore - make_openai_client, - require_llm_env, - resolve_chat_credentials, - ) - from ._compat import record_usage # type: ignore - - require_llm_env(context=context) - key, base_url = resolve_chat_credentials( - model=model, - api_key_env=api_key_env, - base_url_env=base_url_env, - ) - if not key: - raise RuntimeError( - f"{context}: no LLM credentials after resolve_chat_credentials " - f"(model={model!r}; need DS_KEY for deepseek-* or OPENAI_API_KEY)." - ) - client = make_openai_client(api_key=key, base_url=base_url, timeout=timeout) - try: - cached = cached_chat_completion( - client, - purpose=purpose, - model=model, - messages=list(messages), - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - ) - finally: - _logger.info( - "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", - purpose, - thinking_role, - time.perf_counter() - call_started, - ) - if usage_tag: - record_usage(usage_tag, cached.get("usage")) - record_episode_tokens(cached.get("usage")) - return cached diff --git a/deprecated/mapnav/nav/nav_map_scores.py b/deprecated/mapnav/nav/nav_map_scores.py deleted file mode 100644 index 8877451a0..000000000 --- a/deprecated/mapnav/nav/nav_map_scores.py +++ /dev/null @@ -1,333 +0,0 @@ -from __future__ import annotations - -import logging -import time -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple - -from shared.services.retrieval.scoring.knowhere_hybrid import score_persisted_corpus_many -from shared.services.retrieval.scoring.score_units import _walk_tree - -_logger = logging.getLogger(__name__) - - -def _count_tree_shape( - tree_by_doc: Dict[ - str, - Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], - ], -) -> Tuple[int, int, int]: - """Return reachable section nodes, parent-child edges, and leaves.""" - section_nodes: int = sum(len(value[0]) for value in tree_by_doc.values()) - section_edges: int = sum( - sum(len(children) for children in value[0].values()) - for value in tree_by_doc.values() - ) - leaf_sections: int = sum(len(value[1]) for value in tree_by_doc.values()) - return section_nodes, section_edges, leaf_sections - -def _collect_descendant_leaves( - section_id: str, - children_map: Dict[str, List[str]], - leaves: Set[str], -) -> List[str]: - out: List[str] = [] - - def rec(sid: str) -> None: - kids = children_map.get(sid) or [] - if not kids: - if sid in leaves: - out.append(sid) - return - for kid in kids: - rec(kid) - - rec(section_id) - return out - - -def _pool_unit_scores_to_tree( - children_map: Dict[str, List[str]], - leaves: Set[str], - unit_scores: Dict[str, float], -) -> Dict[str, float]: - """MAX-pool globally comparable unit scores onto one document tree.""" - # ``_walk_tree`` inserts every parent before its children. Reversing that - # order is therefore a postorder traversal without allocating descendant - # lists or revisiting nodes. - map_scores: Dict[str, float] = {} - descendant_leaf_max: Dict[str, float] = {} - for section_id in reversed(children_map): - kids = children_map.get(section_id) or [] - if not kids: - leaf_score = float(unit_scores.get(section_id, 0.0) or 0.0) - descendant_leaf_max[section_id] = leaf_score - score = leaf_score - else: - child_leaf_scores = [ - descendant_leaf_max.get(kid, float(unit_scores.get(kid, 0.0) or 0.0)) - for kid in kids - ] - leaf_score = max(child_leaf_scores, default=0.0) - descendant_leaf_max[section_id] = leaf_score - self_score = float(unit_scores.get(f"{section_id}__self", 0.0) or 0.0) - score = max(leaf_score, self_score) - map_scores[section_id] = float(score) - # Preserve the legacy behavior for leaf ids that are not present in the - # children map (defensive support for sparse providers). - for leaf_id in leaves: - map_scores.setdefault(leaf_id, float(unit_scores.get(leaf_id, 0.0) or 0.0)) - return map_scores - - -def compute_map_scores( - ts: Any, - *, - doc_id: str, - query: str, - root_ids: Optional[Sequence[str]] = None, -) -> Dict[str, float]: - """Leaf path+content scores + parent max-pool (self_only only if interstitial).""" - map_scores, _unit_scores = compute_map_and_unit_scores( - ts, doc_id=doc_id, query=query, root_ids=root_ids - ) - return map_scores - - -def compute_map_and_unit_scores( - ts: Any, - *, - doc_id: str, - query: str, - root_ids: Optional[Sequence[str]] = None, - namespace: Optional[str] = None, -) -> Tuple[Dict[str, float], Dict[str, float]]: - """Return (section map_scores, unit hybrid scores keyed by chunk_id).""" - del root_ids - return compute_corpus_map_and_unit_scores( - ts, doc_ids=[doc_id], query=query, namespace=namespace - ) - - -def compute_corpus_map_and_unit_scores( - ts: Any, - *, - doc_ids: Sequence[str], - query: str, - namespace: Optional[str] = None, -) -> Tuple[Dict[str, float], Dict[str, float]]: - """Globally score every unit across documents, then MAX-pool onto the tree. - - All documents share one BM25 corpus, path/content normalization, channel - ranking, and RRF pass. Document-level scores are keyed by bare ``document_id``. - """ - return compute_corpus_map_and_unit_scores_many( - ts, - doc_ids=doc_ids, - queries=[query], - namespace=namespace, - ).get(query, ({}, {})) - - -def compute_corpus_map_and_unit_scores_many( - ts: Any, - *, - doc_ids: Sequence[str], - queries: Sequence[str], - namespace: Optional[str] = None, -) -> Dict[str, Tuple[Dict[str, float], Dict[str, float]]]: - """Globally score several queries with one replay of the corpus units.""" - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - - valid_doc_ids: List[str] = [] - seen_doc_ids: Set[str] = set() - for raw in doc_ids: - doc_id = str(raw or "").strip() - if not doc_id or doc_id in seen_doc_ids: - continue - seen_doc_ids.add(doc_id) - valid_doc_ids.append(doc_id) - - del namespace # Dense scoring is intentionally disabled for the corpus path. - - # Tree shape is query-independent; reuse it across the episode's two - # scoring passes (user query + per-subgoal retrieval_query) on the same - # ToolSpace instead of re-walking every document each time. - tree_cache = getattr(ts, "_mapnav_tree_cache", None) - if not isinstance(tree_cache, dict): - tree_cache = {} - setattr(ts, "_mapnav_tree_cache", tree_cache) - tree_by_doc: Dict[ - str, - Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], - ] = {} - tree_started = time.perf_counter() - for doc_id in valid_doc_ids: - cached = tree_cache.get(doc_id) - if cached is None: - root_ids = list(ts.sections_for_doc(doc_id)) - cached = _walk_tree(ts, doc_id, root_ids) - tree_cache[doc_id] = cached - tree_by_doc[doc_id] = cached - section_nodes, section_edges, leaf_sections = _count_tree_shape(tree_by_doc) - _logger.info( - "retrieval mapnav phase=tree_build seconds=%.3f documents=%d " - "section_nodes=%d section_edges=%d leaf_sections=%d", - time.perf_counter() - tree_started, - len(valid_doc_ids), - section_nodes, - section_edges, - leaf_sections, - ) - - persisted_loader = getattr(ts, "load_persisted_score_corpus", None) - loader_started = time.perf_counter() - persisted_corpus = ( - persisted_loader(valid_doc_ids, unique_queries) - if callable(persisted_loader) - else None - ) - _logger.info( - "retrieval mapnav phase=index_load seconds=%.3f persisted=%s", - time.perf_counter() - loader_started, - persisted_corpus is not None, - ) - score_started = time.perf_counter() - unit_scores_by_query = ( - score_persisted_corpus_many(persisted_corpus, unique_queries) - if persisted_corpus is not None - else {query: {} for query in unique_queries} - ) - _logger.info( - "retrieval mapnav phase=unit_scoring persisted=%s seconds=%.3f units=%d queries=%d", - persisted_corpus is not None, - time.perf_counter() - score_started, - sum(len(scores) for scores in unit_scores_by_query.values()), - len(unique_queries), - ) - results: Dict[str, Tuple[Dict[str, float], Dict[str, float]]] = {} - pooling_started = time.perf_counter() - for query in unique_queries: - unit_scores = unit_scores_by_query.get(query, {}) - map_scores: Dict[str, float] = {} - for doc_id in valid_doc_ids: - children_map, leaves, _titles = tree_by_doc[doc_id] - doc_map_scores = _pool_unit_scores_to_tree( - children_map, leaves, unit_scores - ) - map_scores.update(doc_map_scores) - doc_max = max( - (float(value) for value in doc_map_scores.values()), - default=0.0, - ) - map_scores[doc_id] = doc_max - results[query] = (map_scores, unit_scores) - _logger.info( - "retrieval mapnav phase=map_pooling seconds=%.3f documents=%d " - "section_nodes=%d section_edges=%d leaf_sections=%d", - time.perf_counter() - pooling_started, - len(valid_doc_ids), - section_nodes, - section_edges, - leaf_sections, - ) - return results - - -def unit_id_to_section_id(unit_id: str) -> str: - """Map scoring unit id (leaf or `{sid}__self`) to the section on the map.""" - uid = str(unit_id or "").strip() - if uid.endswith("__self"): - return uid[: -len("__self")] - return uid - - -def select_map_highlights(unit_scores: Dict[str, float], k: int = 6) -> List[str]: - """TOP-K section ids by unit hybrid score (stable tie-break on unit id).""" - limit = max(0, int(k)) - if limit <= 0 or not unit_scores: - return [] - ranked = sorted( - unit_scores.items(), - key=lambda kv: (-float(kv[1] or 0.0), str(kv[0])), - ) - out: List[str] = [] - seen: Set[str] = set() - for uid, _score in ranked: - sid = unit_id_to_section_id(str(uid)) - if not sid or sid in seen: - continue - seen.add(sid) - out.append(sid) - if len(out) >= limit: - break - return out - - -def relight_map_for_query( - ts: Any, - *, - doc_id: str, - query: str, - top_k: int = 6, -) -> Tuple[Dict[str, float], Dict[str, float], List[str]]: - """Re-score the whole shared map against ``query``. - - An empty ``doc_id`` means the corpus root, where document ids are map nodes - and ``ts.document_ids()`` is already restricted to the episode's corpus. - """ - doc = str(doc_id or "").strip() - if doc: - map_scores, unit_scores = compute_map_and_unit_scores( - ts, doc_id=doc, query=query - ) - else: - doc_ids = [str(d) for d in (ts.document_ids() or ()) if str(d).strip()] - if not doc_ids: - return {}, {}, [] - map_scores, unit_scores = compute_corpus_map_and_unit_scores( - ts, doc_ids=doc_ids, query=query - ) - return map_scores, unit_scores, select_map_highlights(unit_scores, k=int(top_k)) - - -def relight_maps_for_queries( - ts: Any, - *, - doc_id: str, - queries: Sequence[str], - top_k: int = 6, -) -> Dict[str, Tuple[Dict[str, float], Dict[str, float], List[str]]]: - """Re-score a shared map for several queries with one corpus replay.""" - unique_queries = list(dict.fromkeys(str(query) for query in queries)) - if not unique_queries: - return {} - doc = str(doc_id or "").strip() - if doc: - return { - query: relight_map_for_query( - ts, - doc_id=doc, - query=query, - top_k=top_k, - ) - for query in unique_queries - } - - doc_ids = [str(value) for value in (ts.document_ids() or ()) if str(value).strip()] - if not doc_ids: - return {} - scored = compute_corpus_map_and_unit_scores_many( - ts, - doc_ids=doc_ids, - queries=unique_queries, - ) - return { - query: ( - map_scores, - unit_scores, - select_map_highlights(unit_scores, k=int(top_k)), - ) - for query, (map_scores, unit_scores) in scored.items() - } diff --git a/deprecated/mapnav/nav/nav_navigate.py b/deprecated/mapnav/nav/nav_navigate.py deleted file mode 100644 index 924cd920d..000000000 --- a/deprecated/mapnav/nav/nav_navigate.py +++ /dev/null @@ -1,109 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, List - -from ._compat import ToolSpace -from .nav_address import owner_document -from .nav_types import LegalAction, NavConfig, NavState - - -def _batch_actions(chosen: LegalAction) -> List[LegalAction]: - batch = (chosen.metadata or {}).get("batch_actions") - if isinstance(batch, list) and batch: - return [a for a in batch if isinstance(a, LegalAction)] - return [chosen] - - -def _section_ancestor_depth(ts: ToolSpace, state: NavState, section_id: str) -> int: - """Structural depth ≈ ancestor count (deeper nodes sort first for COLLECT).""" - sid = str(section_id or "").strip() - if not sid: - return 0 - doc = owner_document(ts, sid, "") or str(state.doc_id or "") - relations = getattr(ts, "section_relation_ids", None) - if not callable(relations): - relations = getattr(ts, "relations", None) - if callable(relations) and doc: - try: - ancestors, _desc = relations(sid, doc) - return len(ancestors or ()) - except Exception: - return 0 - return 0 - - -def _batch_collect_deepest_first( - ts: ToolSpace, state: NavState, chosen: LegalAction -) -> List[LegalAction]: - """Order a COLLECT batch deepest-first so parent purge after child hydrate is stable.""" - acts = [ - a - for a in _batch_actions(chosen) - if str(getattr(a, "section_id", "") or "").strip() - ] - acts.sort( - key=lambda a: ( - -_section_ancestor_depth(ts, state, str(a.section_id)), - str(a.section_id), - ) - ) - return acts - - -def _apply_collect( - ts: ToolSpace, - state: NavState, - chosen: LegalAction, - config: NavConfig, -) -> Dict[str, Any]: - """Run one (possibly batched) COLLECT; mutates state.""" - from .nav_agent import ( # late import avoids cycle - _add_scored, - _collect_subtree, - _mark_collected_branch, - _purge_descendant_evidence, - ) - from .nav_compose import evidence_owner_section_id - - detail: Dict[str, Any] = { - "kind": "collect", - "section_id": chosen.section_id, - "collect_section_ids": [], - "n_added": 0, - "n_hits": 0, - "n_purged_descendant_evidence": 0, - } - total_added = 0 - total_hits = 0 - total_purged = 0 - sids: List[str] = [] - conf_by_sid = dict((chosen.metadata or {}).get("confidence_by_section") or {}) - for act in _batch_collect_deepest_first(ts, state, chosen): - sid = str(act.section_id or "").strip() - if not sid: - continue - # Explicit COLLECT root: record target + confidence from LLM (missing => 0). - state.explicit_collect_ids.add(sid) - conf = float(conf_by_sid.get(sid, 0.0) or 0.0) - state.collect_confidence[sid] = max(0.0, min(1.0, conf)) - # Absorb prior child COLLECTs (e.g. L93) before parent hydrate (L92). - total_purged += _purge_descendant_evidence(ts, state, sid) - scored = _collect_subtree(ts, act, state, config) - # Hydration descendants: confidence stays 0 unless previously explicit. - for chunk, _score in scored: - owner = evidence_owner_section_id(chunk) - if owner and owner != sid: - state.collect_confidence.setdefault(owner, 0.0) - added = _add_scored(state, scored) - cov = _mark_collected_branch(ts, act, state, added) - total_added += added - total_hits += len(scored) - sids.append(sid) - detail.update(cov) - detail["n_added"] = total_added - detail["n_hits"] = total_hits - detail["n_purged_descendant_evidence"] = total_purged - detail["collect_section_ids"] = sids - if sids: - detail["section_id"] = sids[0] - return detail diff --git a/deprecated/mapnav/nav/nav_node_filter.py b/deprecated/mapnav/nav/nav_node_filter.py deleted file mode 100644 index 5e895ee28..000000000 --- a/deprecated/mapnav/nav/nav_node_filter.py +++ /dev/null @@ -1,184 +0,0 @@ -"""Deterministic WHERE filter over the in-memory map-nav hierarchy. - -Agent-authored predicates run on ``path`` (filename + title chain via -``path_titles``) and ``summary``. Field predicates AND together; terms inside -one field OR together. No top-K and no result truncation. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, Iterable, List, Sequence, Tuple - -from shared.services.retrieval.scoring.node_filter_predicates import ( - FieldPredicate, - _compile_predicates, - _node_matches, -) - - -@dataclass(frozen=True) -class NodeFilter: - predicates: Tuple[FieldPredicate, ...] = () - - -@dataclass -class FilterResult: - matched_section_ids: List[str] - matched_doc_ids: List[str] - cardinality: int - failed_predicates: List[str] = field(default_factory=list) - -def node_filter(predicates: Sequence[FieldPredicate]) -> NodeFilter: - return NodeFilter(predicates=tuple(predicates)) - -def apply_node_filter( - ts: Any, - doc_ids: Sequence[str], - nf: NodeFilter, -) -> FilterResult: - """Walk the named documents and evaluate ``nf`` on every node.""" - wanted = [str(did).strip() for did in doc_ids if str(did).strip()] - compiled, failed = _compile_predicates(nf.predicates) - if failed: - return FilterResult( - matched_section_ids=[], - matched_doc_ids=[], - cardinality=0, - failed_predicates=failed, - ) - summaries = _load_summaries(ts) - matched_sections: List[str] = [] - matched_docs: List[str] = [] - seen_sections: set[str] = set() - seen_docs: set[str] = set() - - for doc_id in wanted: - for sid, owner_doc, is_doc_node in _iter_doc_nodes(ts, doc_id): - path_text = _path_text(ts, sid, owner_doc) - summary_text = "" if is_doc_node else str(summaries.get(sid) or "") - if not is_doc_node and not summary_text: - summary_text = _summary_fallback(ts, sid) - values = {"path": path_text, "summary": summary_text} - if not _node_matches(values, compiled): - continue - if is_doc_node: - if owner_doc not in seen_docs: - seen_docs.add(owner_doc) - matched_docs.append(owner_doc) - continue - if sid in seen_sections: - continue - seen_sections.add(sid) - matched_sections.append(sid) - if owner_doc and owner_doc not in seen_docs: - seen_docs.add(owner_doc) - matched_docs.append(owner_doc) - - return FilterResult( - matched_section_ids=matched_sections, - matched_doc_ids=matched_docs, - cardinality=len(matched_sections), - failed_predicates=failed, - ) - - -def render_submap_observation( - ts: Any, - result: FilterResult, - *, - doc_ids: Sequence[str] | None = None, -) -> str: - """Hit-count line plus every matched node (path + summary).""" - del doc_ids - header = f"hits={result.cardinality}" - if result.failed_predicates: - header = f"{header} failed_predicates={len(result.failed_predicates)}" - if result.cardinality == 0: - return header - - summaries = _load_summaries(ts) - lines = [header] - for sid in result.matched_section_ids: - owner = _owner_document(ts, sid) - title = _path_text(ts, sid, owner) or sid - block = [f"{title}"] - summary = str(summaries.get(sid) or "").strip() - if summary: - block.append(f" summary: {summary}") - lines.append("\n".join(block)) - return "\n".join(lines) - -def _iter_doc_nodes(ts: Any, doc_id: str) -> Iterable[Tuple[str, str, bool]]: - yield doc_id, doc_id, True - stack = list(_roots(ts, doc_id)) - seen: set[str] = set() - while stack: - sid = stack.pop(0) - if not sid or sid in seen: - continue - seen.add(sid) - yield sid, doc_id, False - stack[0:0] = [str(child) for child in _children(ts, sid) if str(child)] - - -def _roots(ts: Any, doc_id: str) -> List[str]: - fn = getattr(ts, "sections_for_doc", None) - if callable(fn): - return [str(sid) for sid in (fn(doc_id) or ()) if str(sid).strip()] - provider = getattr(ts, "_provider", None) - root_fn = getattr(provider, "roots", None) if provider is not None else None - if callable(root_fn): - return [str(sid) for sid in (root_fn(doc_id) or ()) if str(sid).strip()] - return [] - - -def _children(ts: Any, section_id: str) -> List[str]: - fn = getattr(ts, "_provider", None) - child_fn = getattr(fn, "children", None) if fn is not None else None - if callable(child_fn): - return [str(sid) for sid in (child_fn(section_id) or ()) if str(sid).strip()] - return [] - - -def _path_text(ts: Any, section_id: str, doc_id: str) -> str: - fn = getattr(ts, "path_titles", None) - if callable(fn): - try: - return str(fn(section_id, doc_id) or "") - except TypeError: - return str(fn(section_id) or "") - return "" - - -def _load_summaries(ts: Any) -> Dict[str, str]: - provider = getattr(ts, "_provider", None) - fn = getattr(provider, "summaries", None) if provider is not None else None - if not callable(fn): - return {} - raw = fn() or {} - return { - str(sid): str(summary or "") - for sid, summary in raw.items() - if str(sid).strip() and str(summary or "").strip() - } - - -def _summary_fallback(ts: Any, section_id: str) -> str: - structure_fn = getattr(ts, "get_structure", None) - if callable(structure_fn): - try: - st = structure_fn(section_id) or {} - return str(st.get("summary") or "").strip() - except Exception: - return "" - return "" - - -def _owner_document(ts: Any, section_id: str) -> str: - fn = getattr(ts, "owner_document", None) - if callable(fn): - got = fn(section_id) - if got: - return str(got) - return "" diff --git a/deprecated/mapnav/nav/nav_orchestrate.py b/deprecated/mapnav/nav/nav_orchestrate.py deleted file mode 100644 index 4ffae8428..000000000 --- a/deprecated/mapnav/nav/nav_orchestrate.py +++ /dev/null @@ -1,816 +0,0 @@ -"""M4/M5: wave orchestration over a RetrievalPlan. - -Execution order = dependency DAG ∩ soft prefer_after. Each subgoal runs its own -harvest so evidence attribution stays per-subgoal. Slot values are -extracted only when a later subgoal references them; checklist acceptance is -owned by ``plan_control``. -""" - -from __future__ import annotations - -import logging -import re -import time -from contextlib import contextmanager -from dataclasses import asdict -from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple - -from .nav_token_budget import stamp_step_detail -from .nav_plan import ( - RetrievalPlan, - Subgoal, - bind_slots, - plan_query, - refine_subgoal_query, - unbound_slots, -) -from .nav_types import NavConfig, NavState, SubgoalResult -from .nav_verify import apply_bindings_from_result, build_subgoal_result - -_SLOT_STRIP_RE = re.compile(r"\{\{\s*[^}]+\s*\}\}") -_logger = logging.getLogger(__name__) - - -def ready_subgoal_ids( - plan: RetrievalPlan, - *, - satisfied: Set[str], - attempted: Optional[Set[str]] = None, - dropped: Optional[Set[str]] = None, -) -> List[str]: - """Subgoals eligible to run now (deps settled, not yet finished). - - F1: a dependency only needs to be *settled* (``satisfied`` or ``dropped``). - Widen leaves a subgoal out of ``attempted`` so it stays ready next wave. - """ - settled = set(satisfied) | set(dropped or ()) - known = {s.id for s in plan.subgoals} - done = set(attempted or ()) | set(satisfied) | set(dropped or ()) - out: List[str] = [] - for sg in plan.subgoals: - if sg.id in done: - continue - deps = [d for d in (sg.depends_on or []) if d in known] - if any(d not in settled for d in deps): - continue - out.append(sg.id) - return order_ready_by_prefer_after(plan, out) - - -def order_ready_by_prefer_after( - plan: RetrievalPlan, - ready_ids: Sequence[str], -) -> List[str]: - """Soft ordering: if A prefer_after B and both ready, B before A.""" - ready = [str(x) for x in ready_ids] - if len(ready) <= 1: - return ready - by = {s.id: s for s in plan.subgoals} - ready_set = set(ready) - # Kahn over soft edges B -> A when A prefer_after B. - indeg = {i: 0 for i in ready} - edges: Dict[str, List[str]] = {i: [] for i in ready} - for sid in ready: - sg = by.get(sid) - if sg is None: - continue - for pred in sg.prefer_after or []: - if pred in ready_set and pred != sid: - edges[pred].append(sid) - indeg[sid] += 1 - queue = [i for i in ready if indeg[i] == 0] - ordered: List[str] = [] - seen = set() - while queue: - # Stable: keep original relative order among zero-indegree nodes. - queue.sort(key=lambda x: ready.index(x)) - node = queue.pop(0) - if node in seen: - continue - seen.add(node) - ordered.append(node) - for nxt in edges.get(node, []): - indeg[nxt] -= 1 - if indeg[nxt] == 0: - queue.append(nxt) - for sid in ready: - if sid not in seen: - ordered.append(sid) - return ordered - - -def _set_focus(state: NavState, subgoal: Subgoal, retrieval_query: str) -> None: - state.focus_subgoal_id = subgoal.id - state.focus_subgoal_need = subgoal.need or retrieval_query - state.focus_retrieval_query = retrieval_query - kind = subgoal.contract.kind - card = subgoal.contract.cardinality - state.focus_contract_kind = str(kind or "") - state.focus_subgoal_contract = ( - f"{kind}" + (f" cardinality={card}" if card is not None else "") - ) - - -def _clear_focus(state: NavState) -> None: - state.focus_subgoal_id = "" - state.focus_subgoal_need = "" - state.focus_subgoal_contract = "" - state.focus_retrieval_query = "" - state.focus_contract_kind = "" - - -def _unbound_retrieval_query(subgoal: Subgoal) -> str: - """Drop unresolved slot braces for REBIND degrade.""" - raw = _SLOT_STRIP_RE.sub(" ", subgoal.retrieval_query or "").strip() - raw = re.sub(r"\s+", " ", raw).strip() - return raw or (subgoal.need or "").strip() or subgoal.retrieval_query - - -def _resolve_subgoal_query(state: NavState, subgoal: Subgoal) -> str: - query = bind_slots(subgoal.retrieval_query, state.slot_bindings) - if unbound_slots(query): - query = _unbound_retrieval_query(subgoal) - refined = str( - (state.subgoal_refined_queries or {}).get(subgoal.id) or "" - ).strip() - return refined or query - - -def seed_episode_map_scores_from_plan( - ts: Any, - state: NavState, - config: NavConfig, - plan: RetrievalPlan, -) -> int: - """Light the map from executable plan retrieval_query strings; set baseline. - - Scores bound (or refined) retrieval_query values only — queries still holding - ``{{slot}}`` placeholders are skipped and scored later when harvest resolves - them. Empty plan falls back to ``state.query``. Harvest may temporarily - relight per subgoal; ``_relit_map`` restores this baseline for evidence pack. - Returns the number of unique queries scored. - """ - queries: List[str] = [] - seen: Set[str] = set() - for subgoal in plan.subgoals: - refined = str( - (state.subgoal_refined_queries or {}).get(subgoal.id) or "" - ).strip() - if refined: - q = refined - else: - bound = bind_slots(subgoal.retrieval_query, state.slot_bindings) - # Skip not-yet-bound slot queries — scoring a stripped skeleton - # cannot cache-hit the later bound harvest string. - if unbound_slots(bound): - continue - q = str(bound or "").strip() - if q and q not in seen: - seen.add(q) - queries.append(q) - if not queries: - q = str(state.query or "").strip() - if q: - queries = [q] - if not queries: - state.map_scores = {} - state.unit_scores = {} - state.highlight_ids = [] - state.relit_map_cache = {} - return 0 - - # Replace prior episode lighting entirely (including after replan). - state.relit_map_cache = {} - - map_started = time.perf_counter() - from .nav_map_scores import relight_maps_for_queries - - prepared = relight_maps_for_queries( - ts, - doc_id=state.doc_id, - queries=queries, - top_k=int(config.collect_top_k), - ) - for query, triple in prepared.items(): - scores, units, highlights = triple - if scores: - state.relit_map_cache[query] = ( - dict(scores), - dict(units), - list(highlights), - ) - - baseline = None - for query in queries: - triple = prepared.get(query) - if triple and triple[0]: - baseline = triple - break - if baseline is not None: - scores, units, highlights = baseline - state.map_scores = dict(scores) - state.unit_scores = dict(units) - state.highlight_ids = list(highlights) - else: - state.map_scores = {} - state.unit_scores = {} - state.highlight_ids = [] - - _logger.info( - "retrieval mapnav phase=map_scoring seconds=%.3f queries=%d", - time.perf_counter() - map_started, - len(queries), - ) - return len(queries) - - -def _wave_subgoal_result( - plan: RetrievalPlan, - state: NavState, - config: NavConfig, - subgoal: Subgoal, - *, - retrieval_query: str, - new_chunks: Sequence[Tuple[Any, float]], - collected_before: Set[str], - explicit_before: Set[str], - steps_out: Optional[List[Any]] = None, -) -> SubgoalResult: - """Build this wave's result from *new* chunks only (never the global pool).""" - return build_subgoal_result( - plan, - state.collected_section_ids, - config, - subgoal, - retrieval_query=retrieval_query, - new_chunks=new_chunks, - collected_before=collected_before, - explicit_collect_ids=state.explicit_collect_ids, - explicit_before=explicit_before, - steps_out=steps_out, - ) - - -@contextmanager -def _relit_map( - ts: Any, - state: NavState, - config: NavConfig, - *, - query: str, - prepared: Optional[Tuple[Dict[str, float], Dict[str, float], List[str]]] = None, -) -> Iterator[None]: - """Score the shared map against the harvest ``query`` for one call. - - Episode-level ``state.map_scores`` is the post-plan baseline (lit from - plan ``retrieval_query`` strings). Checklist harvests may run under a - different per-subgoal string, so the map is re-scored for that call and - restored to the episode baseline afterward for evidence packing. - """ - relit = prepared - q = (query or "").strip() - if relit is None and q: - relit = state.relit_map_cache.get(q) - if relit is None and q: - try: - from .nav_map_scores import relight_map_for_query - - scores, units, highlights = relight_map_for_query( - ts, - doc_id=state.doc_id, - query=q, - top_k=int(config.collect_top_k), - ) - if scores: - relit = (dict(scores), dict(units), list(highlights)) - state.relit_map_cache[q] = relit - except Exception: - relit = None - if relit is None: - yield - return - scores, units, highlights = relit - saved = (state.map_scores, state.unit_scores, state.highlight_ids) - # Copy so harvest readers cannot mutate the cached baseline triple. - state.map_scores = dict(scores) - state.unit_scores = dict(units) - state.highlight_ids = list(highlights) - try: - yield - finally: - state.map_scores, state.unit_scores, state.highlight_ids = saved - - -def _execute_subgoal_harvest_once( - ts: Any, - state: NavState, - config: NavConfig, - plan: RetrievalPlan, - subgoal: Subgoal, - *, - steps_out: Optional[List[Any]], - retrieval_query: Optional[str] = None, - prepared_relight: Optional[ - Tuple[Dict[str, float], Dict[str, float], List[str]] - ] = None, -) -> Dict[str, Any]: - """One harvest() call for this subgoal this wave — no internal retry loop. - - Retry / widen / drop / replan authority belongs to ``plan_control`` across - waves (see ``nav_control.plan_control``), not to this single call. - """ - rq = retrieval_query or _resolve_subgoal_query(state, subgoal) - refined = str((state.subgoal_refined_queries or {}).get(subgoal.id) or "").strip() - _set_focus(state, subgoal, rq) - # Always enter at namespace/document root; prior dead-ends stay hidden via - # subgoal_dismissed_section_ids so the next harvest sees siblings instead. - before_sections = set(state.collected_section_ids) - before_explicit = set(state.explicit_collect_ids) - before_len = len(state.collected) - with _relit_map( - ts, - state, - config, - query=rq, - prepared=prepared_relight, - ): - harvest_result = _harvest_after_node_filter( - ts, - state, - config, - subgoal=subgoal, - query=rq, - steps_out=steps_out, - ) - new_chunks = list(state.collected[before_len:]) - signal = _wave_subgoal_result( - plan, - state, - config, - subgoal, - retrieval_query=rq, - new_chunks=new_chunks, - collected_before=before_sections, - explicit_before=before_explicit, - steps_out=steps_out, - ) - _clear_focus(state) - return { - "subgoal_id": subgoal.id, - "result": signal, - "new_chunks": new_chunks, - "harvest": { - "n_policy_calls": harvest_result.n_policy_calls, - "visited_section_ids": list(harvest_result.visited_section_ids), - "max_depth_hit": harvest_result.max_depth_hit, - "reason": harvest_result.reason, - "widen_gap": str((state.subgoal_widen_gaps or {}).get(subgoal.id) or ""), - "refined_query": refined, - }, - } - - -def _harvest_after_node_filter( - ts: Any, - state: NavState, - config: NavConfig, - *, - subgoal: Subgoal, - query: str, - steps_out: Optional[List[Any]], -) -> Any: - from .nav_harvest import harvest - - if not bool(getattr(config, "enable_node_filter", False)) or not bool( - getattr(subgoal, "use_node_filter", False) - ): - return harvest( - ts, - state, - config, - subgoal=subgoal, - entry_scope=None, - query=query, - steps_out=steps_out, - ) - - from .nav_scope_filter import run_scope_filter - - doc_ids = list(ts.document_ids() or ()) - if not doc_ids and state.doc_id: - doc_ids = [str(state.doc_id)] - seed = None - raw_seed = list(getattr(subgoal, "node_filter_predicates", None) or []) - if raw_seed: - from .nav_scope_filter import parse_node_filter - - seed = parse_node_filter({"predicates": raw_seed}) - outcome = run_scope_filter( - ts, - config, - query=query, - doc_ids=doc_ids, - seed_filter=seed, - steps_out=steps_out, - ) - if outcome.decision == "collect_all": - collected = _collect_filtered_sections( - ts, state, config, subgoal=subgoal, section_ids=outcome.settled_section_ids - ) - if collected is not None: - return collected - if outcome.decision == "scoped_harvest" and outcome.settled_section_ids: - return harvest( - ts, - state, - config, - subgoal=subgoal, - entry_scope=None, - query=query, - steps_out=steps_out, - allowed_section_ids=set(outcome.settled_section_ids), - ) - return harvest( - ts, - state, - config, - subgoal=subgoal, - entry_scope=None, - query=query, - steps_out=steps_out, - ) - - -def _collect_filtered_sections( - ts: Any, - state: NavState, - config: NavConfig, - *, - subgoal: Subgoal, - section_ids: Sequence[str], -) -> Any: - from .nav_address import is_dispatch_only_node - from .nav_harvest import HarvestResult - from .nav_knowhere import is_root_section - from .nav_navigate import _apply_collect - from .nav_types import ActionKind, LegalAction - - collectable = [ - sid - for sid in section_ids - if str(sid).strip() - and not is_dispatch_only_node(ts, sid) - and not is_root_section(ts, sid) - ] - if not collectable: - return None - actions = [ - LegalAction( - action_id=f"C{i}", - kind=ActionKind.COLLECT, - section_id=sid, - metadata={"multi": True}, - ) - for i, sid in enumerate(collectable, start=1) - ] - primary = actions[0] - primary.metadata = dict(primary.metadata or {}) - primary.metadata["batch_actions"] = actions - primary.metadata["confidence_by_section"] = {sid: 1.0 for sid in collectable} - detail = _apply_collect(ts, state, primary, config) - new_roots = list(detail.get("collect_section_ids") or []) - if bool(config.is_checklist): - for sid in new_roots: - state.harvested_owner_subgoal[sid] = subgoal.id - return HarvestResult( - subgoal_id=subgoal.id, - new_section_ids=new_roots, - n_policy_calls=0, - reason="node_filter_collect_all", - ) - - -def _apply_plan_control( - ts: Any, - state: NavState, - config: NavConfig, - *, - plan: RetrievalPlan, - outputs: Sequence[Dict[str, Any]], - by_id: Dict[str, Subgoal], - steps_out: Optional[List[Any]], -) -> Dict[str, Any]: - """Apply one wave's plan_control decision; returns a TRACE-friendly detail. - - ``widen`` = leave the subgoal unsettled for the next wave. plan_control says - what is missing; PLAN turns that into a new retrieval_query, which then - re-scores the shared map for the next harvest. Prior dead-ends stay - dismissed. Only ``subgoal_max_attempts`` turns widen into drop. - """ - from .nav_control import plan_control - - decision = plan_control(ts, state, config, plan=plan, wave_outputs=outputs) - max_attempts = max(1, int(getattr(config, "subgoal_max_attempts", 2) or 2)) - - for item in outputs: - sid = item["subgoal_id"] - result: SubgoalResult = item["result"] - sub_decision = decision.per_subgoal.get(sid) - has_evidence = int(getattr(result, "chars_used", 0) or 0) > 0 - kind = sub_decision.decision if sub_decision else ("accept" if has_evidence else "widen") - # Circuit breaker: bound widen loops regardless of plan_control. - if kind == "widen" and int(state.subgoal_attempt_counts.get(sid, 0)) >= max_attempts: - kind = "drop" - - if kind == "accept": - result.satisfied = True - state.subgoal_results[sid] = asdict(result) - state.satisfied_subgoal_ids.add(sid) - state.attempted_subgoal_ids.add(sid) - state.subgoal_widen_gaps.pop(sid, None) - state.subgoal_refined_queries.pop(sid, None) - elif kind == "drop": - state.dropped_subgoal_ids.add(sid) - state.attempted_subgoal_ids.add(sid) - state.subgoal_widen_gaps.pop(sid, None) - state.subgoal_refined_queries.pop(sid, None) - else: - # widen: control's note says what is missing. It is never fed to - # harvest as text — PLAN reads it against the map and rewrites the - # query, which is what actually re-ranks the next harvest. - note = "" - if sub_decision is not None: - note = str(sub_decision.note or "").strip() - if not note: - note = str(getattr(result, "gap", "") or "").strip() - if note: - state.subgoal_widen_gaps[sid] = note - subgoal = by_id.get(sid) - if subgoal is not None: - refined = refine_subgoal_query( - ts, - state, - config, - subgoal=subgoal, - previous_query=str( - state.subgoal_refined_queries.get(sid) - or subgoal.retrieval_query - or "" - ), - gap=note, - selected_section_ids=list( - getattr(result, "explicit_collect_ids", None) - or getattr(result, "collected_section_ids", None) - or [] - ), - ) - if refined: - state.subgoal_refined_queries[sid] = refined - - if steps_out is not None: - from ._compat import AgentStep # type: ignore - - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="plan_control", - detail=stamp_step_detail({ - "global": decision.global_action, - "reason": decision.reason, - "subgoals": { - sid: {"decision": d.decision, "note": d.note} - for sid, d in decision.per_subgoal.items() - }, - }), - ) - ) - return { - "global": decision.global_action, - "replan": decision.global_action == "replan", - "done": decision.global_action == "done", - "reason": decision.reason, - } - - -def execute_plan( - ts: Any, - state: NavState, - config: NavConfig, - *, - steps_out: Optional[List[Any]] = None, - episode_query: str = "", -) -> Dict[str, Any]: - """Run dependency waves until idle or ``max_waves``.""" - from .nav_agent import AgentStep - - plan = state.retrieval_plan - if plan is None or not getattr(plan, "subgoals", None): - # Retired: no multi-step navigate fallback. Planning always yields a - # subgoal (fallback_plan on failure); an empty plan here is a bug. - _logger.warning("execute_plan called with empty retrieval_plan; skipping wave") - return {"waves": [], "results": {}} - - max_waves = int(getattr(config, "max_waves", 0) or 0) - # Checklist mode always harvests + plan_controls. - wave_idx = 0 - summary: Dict[str, Any] = {"waves": [], "results": {}} - episode_done = False - - while True: - if episode_done: - break - if max_waves > 0 and wave_idx >= max_waves: - break - ready = ready_subgoal_ids( - plan, - satisfied=set(state.satisfied_subgoal_ids), - attempted=set(state.attempted_subgoal_ids), - dropped=set(state.dropped_subgoal_ids), - ) - if not ready: - break - wave_idx += 1 - wave_detail: Dict[str, Any] = { - "wave": wave_idx, - "ready": list(ready), - "subgoal_results": [], - } - - by_id = {s.id: s for s in plan.subgoals} - outputs: List[Dict[str, Any]] = [] - query_by_subgoal = { - sid: str(_resolve_subgoal_query(state, by_id[sid]) or "").strip() - for sid in ready - } - prepared_relights: Dict[ - str, - Tuple[Dict[str, float], Dict[str, float], List[str]], - ] = {} - missing_queries: List[str] = [] - for query in query_by_subgoal.values(): - if not query: - continue - cached = state.relit_map_cache.get(query) - if cached is not None: - prepared_relights[query] = cached - elif query not in missing_queries: - missing_queries.append(query) - if missing_queries: - try: - from .nav_map_scores import relight_maps_for_queries - - scored = relight_maps_for_queries( - ts, - doc_id=state.doc_id, - queries=missing_queries, - top_k=int(config.collect_top_k), - ) - for query, triple in scored.items(): - scores, units, highlights = triple - prepared = (dict(scores), dict(units), list(highlights)) - prepared_relights[query] = prepared - if scores: - state.relit_map_cache[query] = prepared - except Exception: - _logger.exception( - "Failed to precompute relight maps; continuing without prepared relights " - "(doc_id=%s, query_count=%d)", - state.doc_id, - len(missing_queries), - ) - - def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) -> Dict[str, Any]: - query = query_by_subgoal[sid] - prepared = prepared_relights.get(query) - if prepared is not None and not prepared[0]: - prepared = None - return _execute_subgoal_harvest_once( - ts, - working_state, - config, - plan, - by_id[sid], - steps_out=out_steps, - retrieval_query=query, - prepared_relight=prepared, - ) - - # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). - wave_started = time.perf_counter() - for sid in ready: - subgoal_started = time.perf_counter() - try: - outputs.append(_run_one(sid, state, steps_out)) - finally: - _logger.info( - "retrieval mapnav harvest subgoal=%s wave=%d seconds=%.3f", - sid, - wave_idx, - time.perf_counter() - subgoal_started, - ) - - # Bookkeeping shared by both decision paths. - for item in outputs: - sid = item["subgoal_id"] - result: SubgoalResult = item["result"] - state.subgoal_results[sid] = asdict(result) - wave_detail["subgoal_results"].append( - { - "subgoal_id": sid, - "chars_used": result.chars_used, - "gap": result.gap, - "extracted": dict(result.extracted or {}), - } - ) - summary["results"][sid] = asdict(result) - if result.extracted: - state.slot_bindings = apply_bindings_from_result( - state.slot_bindings, by_id[sid], result.extracted - ) - state.subgoal_attempt_counts[sid] = int( - state.subgoal_attempt_counts.get(sid, 0) - ) + 1 - - control_started = time.perf_counter() - try: - control_detail = _apply_plan_control( - ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out - ) - finally: - _logger.info( - "retrieval mapnav plan_control wave=%d seconds=%.3f", - wave_idx, - time.perf_counter() - control_started, - ) - wave_detail["plan_control"] = control_detail - replan_requested = bool(control_detail.get("replan")) - if control_detail.get("done"): - episode_done = True - - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="plan_wave", - detail=stamp_step_detail(wave_detail), - ) - ) - summary["waves"].append(wave_detail) - _logger.info( - "retrieval mapnav wave=%d ready=%d seconds=%.3f", - wave_idx, - len(ready), - time.perf_counter() - wave_started, - ) - - if replan_requested: - cap = int(getattr(config, "max_replans", 0) or 0) - if cap > 0 and int(state.replan_count) < cap: - state.replan_count += 1 - t0 = time.perf_counter() - new_plan = plan_query(state, config) - state.retrieval_plan = new_plan - plan = new_plan - # A regenerated plan gets fresh subgoal ids (s1, s2, ... again), - # so per-id bookkeeping (satisfied/attempted/dropped/widen gaps, - # qualified "sX.slot" bindings) cannot be safely carried over — - # those ids now mean something else. What IS safe and worth - # keeping is unqualified slot bindings (plain fact values) and - # every chunk already in state.collected. Clear before re-seeding - # so stale refined queries / qualified bindings cannot leak into - # the new plan's retrieval_query scoring. - state.satisfied_subgoal_ids = set() - state.attempted_subgoal_ids = set() - state.dropped_subgoal_ids = set() - state.subgoal_results = {} - state.subgoal_widen_gaps = {} - state.subgoal_refined_queries = {} - state.subgoal_attempt_counts = {} - state.subgoal_dismissed_section_ids = {} - state.slot_bindings = { - k: v for k, v in state.slot_bindings.items() if "." not in k - } - seed_episode_map_scores_from_plan(ts, state, config, new_plan) - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="replan", - detail=stamp_step_detail({ - "replan_count": state.replan_count, - "n_subgoals": len(new_plan.subgoals), - "fallback": bool(new_plan.fallback), - "seconds": time.perf_counter() - t0, - }), - ) - ) - continue - # Cap reached: stop requesting further replans. - replan_requested = False - - _clear_focus(state) - summary["n_waves"] = wave_idx - summary["satisfied"] = sorted(state.satisfied_subgoal_ids) - summary["attempted"] = sorted(state.attempted_subgoal_ids) - summary["dropped"] = sorted(state.dropped_subgoal_ids) - return summary diff --git a/deprecated/mapnav/nav/nav_plan.py b/deprecated/mapnav/nav/nav_plan.py deleted file mode 100644 index 6ff770022..000000000 --- a/deprecated/mapnav/nav/nav_plan.py +++ /dev/null @@ -1,941 +0,0 @@ -"""Structure-conditioned query planning (M2). - -``plan_query`` is query-only: it emits a coverage checklist plus an auditable -RetrievalPlan from the user question (no pre-lit map). ``refine_subgoal_query`` -may still read a folded planning map after harvest has lit scores. -""" - -from __future__ import annotations - -import json -import os -import re -from dataclasses import asdict, dataclass, field, replace -from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple - -from .nav_actions import build_legal_actions, format_actionable_map_observation -from .nav_projection import build_map -from .nav_types import NavConfig, NavState, Projection - -ContractKind = Literal[ - "single_fact", - "enumeration", - "span", - "comparison", - "existence", -] -RelationKind = Literal["parent-child", "sibling"] -MapCoverage = Literal["sufficient", "partial", "insufficient"] - -_SLOT_REF_RE = re.compile(r"\{\{\s*([A-Za-z_][\w]*(?:\.[A-Za-z_][\w]*)?)\s*\}\}") -_CONTRACT_KINDS = { - "single_fact", - "enumeration", - "span", - "comparison", - "existence", -} -_RELATION_KINDS = {"parent-child", "sibling"} -_MAP_COVERAGE = {"sufficient", "partial", "insufficient"} -_PLANNER_PURPOSE = "nav_query_plan_v4" -_QUERY_REFINE_PURPOSE = "nav_query_refine_v3" - - -@dataclass -class Contract: - kind: ContractKind = "single_fact" - cardinality: Optional[int] = None - - -@dataclass -class Subgoal: - id: str - need: str - retrieval_query: str - depends_on: List[str] = field(default_factory=list) - prefer_after: List[str] = field(default_factory=list) - contract: Contract = field(default_factory=Contract) - produces: List[str] = field(default_factory=list) - use_node_filter: bool = False - node_filter_predicates: List[Dict[str, Any]] = field(default_factory=list) - - -@dataclass -class SubgoalEdge: - source: str - target: str - kind: RelationKind - - -@dataclass -class CoverageItem: - """One fact that the episode's evidence must eventually cover.""" - - id: str - fact: str - - -@dataclass -class RetrievalPlan: - subgoals: List[Subgoal] = field(default_factory=list) - relations: List[SubgoalEdge] = field(default_factory=list) - coverage_checklist: List[CoverageItem] = field(default_factory=list) - reason: str = "" - map_coverage: MapCoverage = "sufficient" - fallback: bool = False - raw: str = "" - - def to_dict(self) -> Dict[str, Any]: - return asdict(self) - - def subgoal_by_id(self) -> Dict[str, Subgoal]: - return {s.id: s for s in self.subgoals} - - -def bind_slots(text: str, bindings: Dict[str, str]) -> str: - """Fill ``{{slot}}`` / ``{{s1.slot}}`` placeholders from extracted bindings.""" - - def repl(m: re.Match[str]) -> str: - key = m.group(1) - if key in bindings: - return str(bindings[key]) - if "." in key: - short = key.split(".", 1)[1] - if short in bindings: - return str(bindings[short]) - return m.group(0) - - return _SLOT_REF_RE.sub(repl, text or "") - - -def unbound_slots(text: str) -> List[str]: - return [m.group(1) for m in _SLOT_REF_RE.finditer(text or "")] - - -def extract_plan_json(text: str) -> Optional[dict]: - """Parse a (possibly nested) JSON object from model output. - - Uses brace balancing so nested plan objects are not truncated to the first - inner ``{...}``. - """ - s = (text or "").strip() - if not s: - return None - if s.startswith("```"): - s = re.sub(r"^```(?:json)?\s*", "", s) - s = re.sub(r"\s*```$", "", s) - s = s.strip() - if s.startswith("{") and s.endswith("}"): - try: - obj = json.loads(s) - return obj if isinstance(obj, dict) else None - except Exception: - pass - start = s.find("{") - if start < 0: - return None - depth = 0 - in_str = False - esc = False - for i in range(start, len(s)): - ch = s[i] - if in_str: - if esc: - esc = False - elif ch == "\\": - esc = True - elif ch == '"': - in_str = False - continue - if ch == '"': - in_str = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - try: - obj = json.loads(s[start : i + 1]) - return obj if isinstance(obj, dict) else None - except Exception: - return None - return None - - -def _as_bool(raw: Any) -> bool: - if isinstance(raw, bool): - return raw - return str(raw or "").strip().lower() in {"1", "true", "yes", "on"} - - -def _parse_node_filter_predicates(raw: Any) -> List[Dict[str, Any]]: - if isinstance(raw, dict): - raw = raw.get("predicates") - if not isinstance(raw, list): - return [] - out: List[Dict[str, Any]] = [] - for item in raw: - if not isinstance(item, dict): - continue - field = str(item.get("field") or "").strip().lower() - if field not in {"path", "summary"}: - continue - terms = item.get("terms") or [] - if isinstance(terms, str): - terms = [terms] - if not isinstance(terms, list): - continue - cleaned = [str(t) for t in terms if str(t)] - if not cleaned: - continue - match = str(item.get("match") or "substring").strip().lower() - if match not in {"substring", "regex"}: - match = "substring" - out.append({"field": field, "terms": cleaned, "match": match}) - return out - - -def _as_str_list(raw: Any) -> List[str]: - if raw is None: - return [] - if isinstance(raw, str): - s = raw.strip() - return [s] if s else [] - if isinstance(raw, list): - out: List[str] = [] - for x in raw: - t = str(x).strip() - if t: - out.append(t) - return out - return [] - - -def _parse_contract(raw: Any) -> Contract: - if not isinstance(raw, dict): - return Contract() - kind = str(raw.get("kind") or "single_fact").strip().lower().replace("-", "_") - if kind not in _CONTRACT_KINDS: - kind = "single_fact" - card = raw.get("cardinality", None) - cardinality: Optional[int] = None - if card is not None and str(card).strip() != "": - try: - cardinality = max(1, int(card)) - except Exception: - cardinality = None - return Contract( - kind=kind, # type: ignore[arg-type] - cardinality=cardinality, - ) - - -def _script_hist(text: str) -> Dict[str, float]: - """Fraction of non-space chars in cjk / latin / other (no language hardcoding).""" - counts = {"cjk": 0, "latin": 0, "other": 0} - for ch in text or "": - if ch.isspace(): - continue - code = ord(ch) - if 0x4E00 <= code <= 0x9FFF or 0x3400 <= code <= 0x4DBF: - counts["cjk"] += 1 - elif ("A" <= ch <= "Z") or ("a" <= ch <= "z"): - counts["latin"] += 1 - else: - counts["other"] += 1 - total = sum(counts.values()) - if total <= 0: - return {"cjk": 0.0, "latin": 0.0, "other": 0.0} - return {k: float(v) / float(total) for k, v in counts.items()} - - -def _dominant_letter_script(hist: Dict[str, float]) -> Optional[str]: - cjk = float(hist.get("cjk") or 0.0) - latin = float(hist.get("latin") or 0.0) - if cjk <= 0.0 and latin <= 0.0: - return None - return "cjk" if cjk >= latin else "latin" - - -def retrieval_query_language_mismatch(retrieval_query: str, reference: str) -> bool: - """True when rq uses another letter script while lacking the map/query script. - - Criterion is absolute absence of the reference's dominant letter script while - another letter script is present — not a numeric threshold, not a fixed language. - """ - rq = (retrieval_query or "").strip() - ref = (reference or "").strip() - if not rq or not ref: - return False - ref_dom = _dominant_letter_script(_script_hist(ref)) - if ref_dom is None: - return False - q_hist = _script_hist(rq) - if ref_dom == "cjk" and q_hist.get("cjk", 0.0) == 0.0 and q_hist.get("latin", 0.0) > 0.0: - return True - if ref_dom == "latin" and q_hist.get("latin", 0.0) == 0.0 and q_hist.get("cjk", 0.0) > 0.0: - return True - return False - - -def language_reference_text( - *, - query: str, - projection: Optional[Projection] = None, -) -> str: - """Script reference for retrieval_query checks. - - Planner is query-only (projection=None). Refine still sees the folded map; - its English chrome (collect=/dispatch=/[Hit]) is excluded — only visible - map *titles* are used so they do not dominate script detection. - """ - parts: List[str] = [] - q = (query or "").strip() - if q: - parts.append(q) - if projection is not None: - for v in projection.tree_sections or projection.visible_sections or []: - title = str(getattr(v, "title", "") or "").strip() - if not title: - title = str(getattr(v, "preview", "") or "").strip() - if title: - parts.append(title) - return "\n".join(parts) - - -def plan_has_language_mismatch(plan: RetrievalPlan, reference: str) -> bool: - return any( - retrieval_query_language_mismatch(s.retrieval_query, reference) - for s in plan.subgoals - ) - - -def _has_cycle(subgoals: List[Subgoal]) -> bool: - deps: Dict[str, List[str]] = {s.id: list(s.depends_on) for s in subgoals} - visiting: Set[str] = set() - done: Set[str] = set() - - def dfs(node: str) -> bool: - if node in done: - return False - if node in visiting: - return True - visiting.add(node) - for nxt in deps.get(node, []): - if nxt in deps and dfs(nxt): - return True - visiting.remove(node) - done.add(node) - return False - - return any(dfs(sid) for sid in deps) - - -def _break_dependency_cycles(subgoals: List[Subgoal]) -> None: - """Move cycle-forming back-edges to prefer_after; keep acyclic depends_on.""" - deps: Dict[str, List[str]] = {s.id: list(s.depends_on) for s in subgoals} - visiting: Set[str] = set() - done: Set[str] = set() - back_edges: List[Tuple[str, str]] = [] - - def dfs(node: str) -> None: - visiting.add(node) - for nxt in list(deps.get(node, [])): - if nxt not in deps: - continue - if nxt in visiting: - back_edges.append((node, nxt)) - elif nxt not in done: - dfs(nxt) - visiting.discard(node) - done.add(node) - - for sid in list(deps.keys()): - if sid not in done: - dfs(sid) - - if not back_edges: - return - by = {s.id: s for s in subgoals} - for src, tgt in back_edges: - s = by.get(src) - if s is None or tgt not in s.depends_on: - continue - s.depends_on = [d for d in s.depends_on if d != tgt] - if tgt not in s.prefer_after: - s.prefer_after.append(tgt) - - -def _slot_owner_and_name(ref: str) -> Tuple[Optional[str], str]: - token = (ref or "").strip() - if not token: - return None, "" - if "." in token: - owner, name = token.split(".", 1) - return owner.strip() or None, name.strip() - return None, token - - -def _apply_slot_dependency_inference(subgoals: List[Subgoal]) -> None: - """Hard deps + produces filled from ``{{s1.slot}}`` references (mechanical).""" - by_id = {s.id: s for s in subgoals} - for s in subgoals: - for ref in unbound_slots(s.retrieval_query) + unbound_slots(s.need): - owner, name = _slot_owner_and_name(ref) - if owner and owner in by_id and owner != s.id: - if owner not in s.depends_on: - s.depends_on.append(owner) - if name and not by_id[owner].produces: - by_id[owner].produces = [name] - - -def _parse_map_coverage(raw: Any) -> MapCoverage: - val = str(raw or "").strip().lower() - if val in _MAP_COVERAGE: - return val # type: ignore[return-value] - return "sufficient" - - -def _parse_coverage_checklist(raw: Any) -> List[CoverageItem]: - if not isinstance(raw, list): - return [] - out: List[CoverageItem] = [] - seen: Set[str] = set() - for i, row in enumerate(raw): - if isinstance(row, str): - fact = row.strip() - cid = f"c{i + 1}" - elif isinstance(row, dict): - fact = str(row.get("fact") or row.get("need") or row.get("item") or "").strip() - cid = str(row.get("id") or f"c{i + 1}").strip() or f"c{i + 1}" - else: - continue - if not fact or cid in seen: - continue - seen.add(cid) - out.append(CoverageItem(id=cid, fact=fact)) - return out - - -def parse_retrieval_plan( - obj: dict, - *, - query: str, -) -> RetrievalPlan: - """Parse LLM JSON into a RetrievalPlan; invalid refs are dropped.""" - rows = obj.get("subgoals") or obj.get("goals") or obj.get("steps") or [] - if not isinstance(rows, list): - rows = [] - - subgoals: List[Subgoal] = [] - seen_ids: Set[str] = set() - for i, row in enumerate(rows): - if not isinstance(row, dict): - continue - sid = str(row.get("id") or row.get("subgoal_id") or f"s{i + 1}").strip() - if not sid or sid in seen_ids: - continue - seen_ids.add(sid) - need = str(row.get("need") or row.get("goal") or "").strip() - rq = str( - row.get("retrieval_query") or row.get("query") or need or query - ).strip() - if not need: - need = rq or query - if not rq: - rq = need or query - produces = _as_str_list(row.get("produces")) - # TODO: revisit one-slot-per-subgoal hard clip (produces[:1]) — - # prompt + validate also enforce it; may want multi-slot later. - if len(produces) > 1: - produces = produces[:1] - subgoals.append( - Subgoal( - id=sid, - need=need, - retrieval_query=rq, - depends_on=_as_str_list(row.get("depends_on")), - prefer_after=_as_str_list(row.get("prefer_after")), - contract=_parse_contract(row.get("contract")), - produces=produces, - use_node_filter=_as_bool(row.get("use_node_filter")), - node_filter_predicates=_parse_node_filter_predicates( - row.get("node_filter") or row.get("node_filter_predicates") - ), - ) - ) - - known = {s.id for s in subgoals} - for s in subgoals: - s.depends_on = [d for d in s.depends_on if d in known and d != s.id] - s.prefer_after = [d for d in s.prefer_after if d in known and d != s.id] - - _apply_slot_dependency_inference(subgoals) - # TODO: same one-slot-per-subgoal policy as parse above; revisit with multi-slot. - for s in subgoals: - if len(s.produces) > 1: - s.produces = s.produces[:1] - - relations: List[SubgoalEdge] = [] - rel_raw = obj.get("relations") or obj.get("edges") or [] - if isinstance(rel_raw, list): - for edge in rel_raw: - if not isinstance(edge, dict): - continue - src = str(edge.get("source") or edge.get("from") or edge.get("a") or "").strip() - tgt = str(edge.get("target") or edge.get("to") or edge.get("b") or "").strip() - kind = str(edge.get("kind") or edge.get("relation") or "").strip().lower() - if kind == "parent_child": - kind = "parent-child" - if kind == "unrelated": - continue - if src not in known or tgt not in known or src == tgt: - continue - if kind not in _RELATION_KINDS: - continue - relations.append(SubgoalEdge(source=src, target=tgt, kind=kind)) # type: ignore[arg-type] - if kind == "parent-child": - for s in subgoals: - if s.id == tgt and src not in s.depends_on: - s.depends_on.append(src) - - if _has_cycle(subgoals): - _break_dependency_cycles(subgoals) - - reason = str(obj.get("reason") or obj.get("plan_reason") or "").strip() - coverage = _parse_map_coverage(obj.get("map_coverage") or obj.get("coverage")) - checklist = _parse_coverage_checklist( - obj.get("coverage_checklist") or obj.get("checklist") - ) - if not checklist: - checklist = [ - CoverageItem(id=f"c{i + 1}", fact=s.need) - for i, s in enumerate(subgoals) - if (s.need or "").strip() - ] - return RetrievalPlan( - subgoals=subgoals, - relations=relations, - coverage_checklist=checklist, - reason=reason, - map_coverage=coverage, - ) - - -def fallback_plan(query: str, *, reason: str = "fallback") -> RetrievalPlan: - """Single shared-space subgoal when planning fails or is unnecessary.""" - q = (query or "").strip() or "query" - return RetrievalPlan( - subgoals=[ - Subgoal( - id="s1", - need=q, - retrieval_query=q, - contract=Contract(kind="single_fact"), - ) - ], - coverage_checklist=[CoverageItem(id="c1", fact=q)], - reason=reason, - map_coverage="sufficient", - fallback=True, - ) - - -def validate_retrieval_plan(plan: RetrievalPlan) -> Tuple[bool, str]: - if not plan.subgoals: - return False, "empty_subgoals" - ids = [s.id for s in plan.subgoals] - if len(ids) != len(set(ids)): - return False, "duplicate_ids" - known = set(ids) - for s in plan.subgoals: - if not (s.need or s.retrieval_query): - return False, f"empty_need:{s.id}" - for d in s.depends_on: - if d not in known: - return False, f"bad_depends_on:{s.id}->{d}" - for d in s.prefer_after: - if d not in known: - return False, f"bad_prefer_after:{s.id}->{d}" - if s.contract.kind not in _CONTRACT_KINDS: - return False, f"bad_contract:{s.id}" - # TODO: revisit multi_produces rejection with one-slot-per-subgoal clip. - if len(s.produces) > 1: - return False, f"multi_produces:{s.id}" - for ref in unbound_slots(s.retrieval_query) + unbound_slots(s.need): - owner, _name = _slot_owner_and_name(ref) - if owner and owner in known and owner not in s.depends_on and owner != s.id: - return False, f"slot_missing_depends:{s.id}->{owner}" - if plan.map_coverage not in _MAP_COVERAGE: - return False, "bad_map_coverage" - if _has_cycle(plan.subgoals): - return False, "cyclic_depends_on" - return True, "ok" - - -def planning_char_limit(config: NavConfig) -> int: - planned = int(getattr(config, "planning_map_char_limit", 0) or 0) - base = int(config.map_char_limit or 0) - if planned > 0: - return planned - return max(1, base) - - -def build_planning_observation( - ts: Any, - state: NavState, - config: NavConfig, -) -> Tuple[Projection, str]: - """Build the one-shot planning map (wider budget; same fold/title-only rules).""" - limit = planning_char_limit(config) - plan_cfg = replace(config, map_char_limit=limit) - projection = build_map( - ts, - doc_id=state.doc_id, - query=state.query, - scope=None, - config=plan_cfg, - map_scores=state.map_scores, - collected_section_ids=state.collected_section_ids, - dismissed_section_ids=state.dismissed_section_ids, - highlight_ids=state.highlight_ids, - harvested_section_ids=( - state.harvested_owner_subgoal if config.is_checklist else None - ), - ) - actions = build_legal_actions( - state, - projection, - config=plan_cfg, - depth=0, - ts=ts, - ) - obs = format_actionable_map_observation( - projection, - actions, - inline_summary=False, - ) - return projection, obs - - -def _planner_system_prompt(*, max_subgoals: int) -> str: - cap = "" - if max_subgoals > 0: - cap = f" Prefer at most {max_subgoals} subgoals." - return ( - "You are a retrieval planner. You receive only the user query (no " - "document map). Emit a coverage checklist plus a retrieval plan over " - "ONE shared search space — do not invent per-subgoal corpus partitions.\n\n" - "Rules:\n" - "1. coverage_checklist lists the facts that episode evidence must cover " - "(short, concrete facts in the query's language)." - f"{cap}\n" - "2. Default to a SINGLE subgoal. Only add more subgoals for a hard data " - "dependency ({{s1.slot}} in a later query) or for clearly independent " - "cross-entity comparisons. Do not split merely to list checklist items.\n" - "3. Each subgoal produces at most ONE slot name in produces " - "(enumeration = one list-valued slot).\n" - "4. retrieval_query is a SHORT KEYWORD QUERY for THIS subgoal only " - "(space-separated entity/role/topic tokens, e.g. \"王仁坤 总工程师 设计成果\"). " - "Split or adapt the user question into compact lexical terms — do NOT " - "emit a full natural-language question or long prose. Downstream lexical " - "retrieval scores and lights the map from these tokens, so keep entity " - "names and role terms; drop filler words. Same language/script as the " - "user query — do not translate terms into another script.\n" - "5. If a later retrieval_query needs a value from an earlier subgoal, " - "write it as {{s1.slot}} (not prose). That implies depends_on.\n" - "6. depends_on = hard data dependency. prefer_after = soft ordering only.\n" - "7. All subgoals share one search space — do not invent per-subgoal scopes.\n" - "8. relations only for parent-child or sibling (omit unrelated pairs).\n" - "9. map_coverage: sufficient | partial | insufficient — whether the user " - "query alone is enough to write executable retrieval_query terms " - "(use sufficient unless the query is empty or unusable).\n" - "10. reason must be English, under 40 words.\n" - "11. Set use_node_filter=true when the subgoal enumerates or compares " - "named facets you can write as path predicates (filenames, section " - "titles). Keep it false for vague semantic needs; " - "retrieval_query remains the fuzzy-leg fallback. Optional node_filter " - "may seed predicates: " - '[{\"field\":\"path\",\"terms\":[\"...\"],\"match\":\"substring|regex\"}].\n\n' - "Return ONLY one JSON object:\n" - "{\n" - ' "reason": "...",\n' - ' "map_coverage": "sufficient|partial|insufficient",\n' - ' "coverage_checklist": [\n' - ' {"id": "c1", "fact": "..."},\n' - ' {"id": "c2", "fact": "..."}\n' - " ],\n" - ' "subgoals": [\n' - " {\n" - ' "id": "s1",\n' - ' "need": "...",\n' - ' "retrieval_query": "...",\n' - ' "depends_on": [],\n' - ' "prefer_after": [],\n' - ' "produces": ["entity"],\n' - ' "contract": {"kind": "single_fact|enumeration|span|comparison|existence", ' - '"cardinality": null},\n' - ' "use_node_filter": false,\n' - ' "node_filter": []\n' - " },\n" - " {\n" - ' "id": "s2",\n' - ' "need": "...",\n' - ' "retrieval_query": "... {{s1.entity}} ...",\n' - ' "depends_on": ["s1"],\n' - ' "produces": ["detail"],\n' - ' "contract": {"kind": "single_fact"}\n' - " }\n" - " ],\n" - ' "relations": [{"source": "s1", "target": "s2", "kind": "parent-child|sibling"}]\n' - "}" - ) - - -def _language_repair_user( - state: NavState, - bad_plan: RetrievalPlan, - *, - reference: str, -) -> str: - bad_lines = [] - for s in bad_plan.subgoals: - if retrieval_query_language_mismatch(s.retrieval_query, reference): - bad_lines.append(f"- {s.id}: {s.retrieval_query}") - bad_block = "\n".join(bad_lines) if bad_lines else "(see prior plan)" - return ( - f"User query: {state.query}\n" - f"Task type: {state.task_type}\n\n" - "Your previous plan had retrieval_query values that do not match the " - "language/script of the user query. Those queries will fail lexical " - "retrieval. Rewrite the FULL plan JSON. Keep structure, but rewrite " - "every mismatched retrieval_query as a short keyword query in the " - "query's own language and terms (space-separated tokens, not a " - "full-sentence question).\n" - f"Mismatched retrieval_query lines:\n{bad_block}\n" - ) - - -def plan_query( - state: NavState, - config: NavConfig, -) -> RetrievalPlan: - """LLM plan from the user query only (no pre-lit map). Falls back to one subgoal.""" - from .nav_llm import ( # type: ignore - nav_chat, - planner_output_max_tokens, - resolve_nav_model, - resolve_nav_thinking_mode, - ) - from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted - - if nav_token_budget_exhausted(): - return fallback_plan(state.query, reason="token_limit") - - model = resolve_nav_model( - model=config.planner_model, - model_env="NAV_PLANNER_MODEL", - fallback_envs=("NAV_LLM_MODEL",), - ) - thinking_mode = resolve_nav_thinking_mode(role="planner") - # Thinking can exceed the default 60s client timeout. - timeout_s = float(os.environ.get("NAV_PLANNER_TIMEOUT_SECONDS", "").strip() or "0") - if timeout_s <= 0: - timeout_s = 300.0 if thinking_mode == "enabled" else 90.0 - max_subgoals = int(getattr(config, "planner_max_subgoals", 0) or 0) - system = _planner_system_prompt(max_subgoals=max_subgoals) - reference = language_reference_text(query=state.query) - user = ( - f"User query: {state.query}\n" - f"Task type: {state.task_type}\n\n" - "Return the retrieval plan JSON." - ) - max_tokens = planner_output_max_tokens( - max( - int(config.llm_max_tokens or 0), - int(getattr(config, "planner_llm_max_tokens", 0) or 0), - 256, - ) - ) - last_raw = "" - last_err = "" - language_repair_used = False - for _attempt in range(3): - try: - cached = nav_chat( - purpose=_PLANNER_PURPOSE, - model=model, - messages=[ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=max_tokens, - response_format={"type": "json_object"}, - thinking_role="planner", - context="Nav Query Planner", - api_key_env="NAV_PLANNER_API_KEY", - base_url_env="NAV_PLANNER_BASE_URL", - timeout=timeout_s, - usage_tag="nav_plan", - ) - content = str(cached.get("content") or "").strip() - reasoning = str(cached.get("reasoning_content") or "").strip() - last_raw = content - if not last_raw and reasoning: - # Only fall back to reasoning if it actually embeds JSON. - if extract_plan_json(reasoning): - last_raw = reasoning - else: - last_err = "empty_content_after_think" - continue - if not last_raw: - last_err = "empty_content" - continue - obj = extract_plan_json(last_raw) or {} - plan = parse_retrieval_plan(obj, query=state.query) - plan.raw = last_raw[:2000] - ok, why = validate_retrieval_plan(plan) - if not ok: - last_err = why - continue - if plan_has_language_mismatch(plan, reference) and not language_repair_used: - language_repair_used = True - last_err = "language_mismatch" - user = _language_repair_user(state, plan, reference=reference) - continue - if plan_has_language_mismatch(plan, reference): - last_err = "language_mismatch" - continue - return plan - except NavTokenLimit: - return fallback_plan(state.query, reason="token_limit") - except Exception as exc: # pragma: no cover - network/LLM path - last_err = str(exc) - continue - plan = fallback_plan(state.query, reason=f"fallback:{last_err or 'unparsed'}") - plan.raw = last_raw[:2000] - return plan - - -def _refine_system_prompt() -> str: - return ( - "You rewrite ONE retrieval_query for a single subgoal on a hierarchical " - "document map.\n" - "A previous attempt for THIS subgoal failed. A controller reported what " - "is still missing. You see the folded title map.\n\n" - "Rules:\n" - "1. Scope is this subgoal only. Rewrite for its need; do not pull in " - "other subgoals' goals or other parts of a multi-part question.\n" - "2. Write a SHORT KEYWORD QUERY (space-separated tokens), not a full " - "sentence. Start from the previous retrieval_query and revise by " - "adding, dropping, or swapping terms so it aims at the reported gap. " - "A paraphrase that ranks the same nodes is useless.\n" - "3. If prior selected nodes are listed, do not aim the question back at " - "those same nodes; steer toward other map regions that may hold the " - "missing evidence, using the map's own wording where helpful.\n" - "4. Same language/script as the map titles. Keep entity names and scope " - "constraints. Do not paste the controller's note verbatim.\n" - "5. Keep it under 40 words, no commentary.\n\n" - 'Return ONLY one JSON object: {"retrieval_query": "..."}' - ) - - -def _section_titles_for_refine(ts: Any, section_ids: Sequence[str]) -> List[str]: - """Titles of previously selected nodes, falling back to the section id.""" - out: List[str] = [] - seen: Set[str] = set() - for raw in section_ids or (): - sid = str(raw or "").strip() - if not sid or sid in seen: - continue - seen.add(sid) - title = "" - try: - st = ts.get_structure(sid) or {} - title = str(st.get("title") or "").strip() - except Exception: - title = "" - out.append(title or sid) - return out - - -def refine_subgoal_query( - ts: Any, - state: NavState, - config: NavConfig, - *, - subgoal: Subgoal, - previous_query: str, - gap: str, - selected_section_ids: Optional[Sequence[str]] = None, -) -> str: - """PLAN-side rewrite of one subgoal's retrieval_query after ``widen``. - - ``plan_control`` can say *what* is missing but never sees the map, so it - cannot write a query that matches the corpus' own terms. This is the one - call that can: it reads the planning map and turns the gap into new search - terms for THIS subgoal only. Returns "" when the rewrite fails, echoes the - previous query, or lands in the wrong script — the caller then keeps the - previous query. - """ - from .nav_llm import nav_chat, resolve_nav_model - from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted - - prev = (previous_query or subgoal.retrieval_query or "").strip() - need = (subgoal.need or prev).strip() - if nav_token_budget_exhausted(): - return "" - - projection, observation = build_planning_observation(ts, state, config) - # Language check against this subgoal + map titles only (no episode query). - reference = language_reference_text(query=need, projection=projection) - selected_titles = _section_titles_for_refine(ts, selected_section_ids or ()) - if selected_titles: - selected_block = "Prior selected nodes:\n" + "\n".join( - f"- {t}" for t in selected_titles - ) - else: - selected_block = "Prior selected nodes: (none)" - model = resolve_nav_model( - model=config.planner_model, - model_env="NAV_PLANNER_MODEL", - fallback_envs=("NAV_LLM_MODEL",), - ) - user = ( - f"Subgoal need: {need}\n" - f"Previous retrieval_query (failed): {prev}\n" - f"Reported gap: {(gap or '').strip() or '(none given)'}\n" - f"{selected_block}\n\n" - f"=== Planning Map ===\n{observation}\n=== End Planning Map ===\n\n" - "Return the rewritten retrieval_query JSON." - ) - try: - cached = nav_chat( - purpose=_QUERY_REFINE_PURPOSE, - model=model, - messages=[ - {"role": "system", "content": _refine_system_prompt()}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=max(128, int(config.llm_max_tokens or 0)), - response_format={"type": "json_object"}, - thinking_role="action", - context="Nav Query Refine", - api_key_env="NAV_PLANNER_API_KEY", - base_url_env="NAV_PLANNER_BASE_URL", - usage_tag="nav_query_refine", - ) - except NavTokenLimit: - return "" - except Exception: # pragma: no cover - network/LLM path - return "" - - obj = extract_plan_json(str(cached.get("content") or "")) or {} - refined = str(obj.get("retrieval_query") or obj.get("query") or "").strip() - if not refined or refined == prev: - return "" - if retrieval_query_language_mismatch(refined, reference): - return "" - return refined diff --git a/deprecated/mapnav/nav/nav_policy.py b/deprecated/mapnav/nav/nav_policy.py deleted file mode 100644 index 666b7ad83..000000000 --- a/deprecated/mapnav/nav/nav_policy.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -import json -import re -from typing import Optional - - -def _extract_json_obj(text: str) -> Optional[dict]: - s = (text or "").strip().replace("```json", "").replace("```", "") - if not s: - return None - if s.startswith("{") and s.endswith("}"): - try: - obj = json.loads(s) - return obj if isinstance(obj, dict) else None - except Exception: - pass - m = re.search(r"\{.*?\}", s, flags=re.S) - if not m: - return None - try: - obj = json.loads(m.group(0)) - return obj if isinstance(obj, dict) else None - except Exception: - return None diff --git a/deprecated/mapnav/nav/nav_projection.py b/deprecated/mapnav/nav/nav_projection.py deleted file mode 100644 index 16b7f2a46..000000000 --- a/deprecated/mapnav/nav/nav_projection.py +++ /dev/null @@ -1,572 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Set - -from .nav_types import NavConfig, Projection, SectionView - -try: - from section_summary_store import get_summary as _store_get_summary -except Exception: # pragma: no cover - def _store_get_summary(section_id: str, *, doc_id: str = "") -> Optional[str]: # type: ignore - return None - - -def _section_summary_for_map( - ts: Any, - section_id: str, - *, - doc_id: str = "", -) -> str: - """Inline map summary: store first, then provider structure (KH production).""" - sid = str(section_id or "").strip() - if not sid: - return "" - try: - text = str(_store_get_summary(sid, doc_id=str(doc_id or "")) or "").strip() - if text: - return text - except Exception: - pass - if ts is None: - return "" - try: - st = ts.get_structure(sid) or {} - text = str(st.get("summary") or "").strip() - if text: - return text - return str(st.get("preview") or "").strip() - except Exception: - return "" - - -def _children(ts: Any, section_id: str) -> List[dict]: - children_fn = getattr(ts, "_children_for_section_path", None) - if callable(children_fn): - loc = getattr(ts, "_idx", None) - doc_id = "" - if loc is not None: - node = getattr(loc, "_node_to_doc_line", {}).get(section_id) - if node: - doc_id = node[0] - else: - synth = getattr(ts, "_synthetic_doc_id", lambda _x: None)(section_id) - doc_id = synth or "" - if doc_id: - rows = children_fn(section_id, doc_id) - if isinstance(rows, list): - return [c for c in rows if isinstance(c, dict)] - st = ts.get_structure(section_id) - children = st.get("children") or [] - if isinstance(children, list): - return [c for c in children if isinstance(c, dict)] - return [] - - -def _top_sections(ts: Any, doc_id: str) -> List[str]: - try: - return list(ts.sections_for_doc(doc_id)) - except Exception: - return [] - - -def _title_from_row(row: dict, section_id: str) -> str: - preview = str(row.get("preview") or "").replace("\n", " ").strip() - if preview: - return preview[:80] - return section_id - - -@dataclass -class _MapNode: - section_id: str - depth: int - title: str - score: float - n_lines: int - has_children: bool - children: List["_MapNode"] = field(default_factory=list) - n_descendants: int = 0 - hidden: bool = False - is_highlight: bool = False - map_id: str = "" - parent_id: Optional[str] = None - harvested_by: str = "" - - -def _count_descendants(node: _MapNode) -> int: - total = 0 - for child in node.children: - if child.hidden: - continue - total += 1 + _count_descendants(child) - node.n_descendants = total - return total - - -def _flatten(nodes: List[_MapNode], *, include_hidden: bool = False) -> List[_MapNode]: - out: List[_MapNode] = [] - - def rec(node: _MapNode) -> None: - if node.hidden and not include_hidden: - return - out.append(node) - for child in node.children: - rec(child) - - for n in nodes: - rec(n) - return out - - -def _ancestors_in_tree(roots: List[_MapNode], section_id: str) -> List[str]: - found: List[str] = [] - - def dfs(node: _MapNode, trail: List[str]) -> bool: - if node.section_id == section_id: - found.extend(trail) - return True - for child in node.children: - if dfs(child, trail + [node.section_id]): - return True - return False - - for root in roots: - if dfs(root, []): - break - return found - - -def _protected_spine_ids(nodes: List[_MapNode], must_keep: Set[str]) -> Set[str]: - """Return must-keep nodes and every ancestor needed to expose them.""" - protected: Set[str] = set() - - def visit(node: _MapNode) -> bool: - contains = node.section_id in must_keep - for child in node.children: - contains = visit(child) or contains - if contains: - protected.add(node.section_id) - return contains - - for root in nodes: - visit(root) - return protected - - -def _mark_hidden_subtree(node: _MapNode) -> None: - node.hidden = True - for child in node.children: - _mark_hidden_subtree(child) - - -def _estimate_actionable_line(node: _MapNode, *, with_summary: bool = False) -> int: - indent = 2 * node.depth - title = len(node.title or node.section_id) - meta = 28 - tags = 14 - actions = 48 - summary = 140 if with_summary else 0 - return indent + title + meta + tags + actions + summary + 1 - - -def _estimate_actionable_total(nodes: List[_MapNode], *, with_summary: bool = False) -> int: - total = 120 - stack = list(nodes) - while stack: - cur = stack.pop() - if cur.hidden: - continue - total += _estimate_actionable_line(cur, with_summary=with_summary) - stack.extend(cur.children) - return total - - -def _visible_subtree_estimate( - node: _MapNode, - *, - with_summary: bool = False, -) -> int: - """Estimated actionable characters removed by hiding this visible subtree.""" - total = 0 - stack = [node] - while stack: - current = stack.pop() - if current.hidden: - continue - total += _estimate_actionable_line( - current, - with_summary=with_summary, - ) - stack.extend(current.children) - return total - - -def _apply_budget_hide( - nodes: List[_MapNode], - *, - char_limit: int, - must_keep: Set[str], - extra_hidden_ids: Optional[Set[str]] = None, - with_summary: bool = False, -) -> None: - """Hide low-score branches so the actionable map fits char_limit. - - TODO: if exploration near highlights fails, re-reveal previously hidden - subtrees into the actionable map (budget permitting). - - Never hard-truncates: keeps hiding score-ordered candidates (depth-0 allowed) - until the estimate fits or only must_keep spine remains. - """ - flat_all = _flatten(nodes, include_hidden=True) - protected = _protected_spine_ids(nodes, must_keep) - extra = set(extra_hidden_ids or ()) - - for node in flat_all: - if node.section_id in extra and node.section_id not in protected: - _mark_hidden_subtree(node) - - for root in nodes: - _count_descendants(root) - - current_estimate = _estimate_actionable_total( - nodes, - with_summary=with_summary, - ) - if current_estimate <= char_limit: - return - - # Sort once. Chosen subtrees are disjoint after already-hidden descendants - # are skipped, so total subtree-estimation work remains linear. - candidates = [ - node - for node in flat_all - if not node.hidden and node.section_id not in protected - ] - candidates.sort( - key=lambda node: ( - node.score, - -node.n_descendants, - -node.depth, - node.section_id, - ) - ) - for node in candidates: - if current_estimate <= char_limit: - break - if node.hidden: - continue - current_estimate -= _visible_subtree_estimate( - node, - with_summary=with_summary, - ) - _mark_hidden_subtree(node) - - -def _build_map_tree( - ts: Any, - *, - root_ids: List[str], - map_scores: Dict[str, float], - collected_section_ids: Optional[Set[str]] = None, - dismissed_section_ids: Optional[Set[str]] = None, - harvested_section_ids: Optional[Dict[str, str]] = None, - keep_ids: Optional[Set[str]] = None, -) -> List[_MapNode]: - roots: List[_MapNode] = [] - seen: Set[str] = set() - harvested = dict(harvested_section_ids or {}) - keep = set(keep_ids) if keep_ids is not None else None - # collected = branch done (sid ∪ descendants already marked by caller). - # Harvested roots stay visible as a collapsed single line (fix-map- - # visibility); their descendants are still fully removed like any other - # collected branch — the point is coverage context, not re-expansion. - gone = ( - set(collected_section_ids or ()) - set(harvested.keys()) - ) | set(dismissed_section_ids or ()) - - def make_node(section_id: str, depth: int, parent_id: Optional[str]) -> Optional[_MapNode]: - if not section_id or section_id in seen: - return None - if section_id in gone: - return None - seen.add(section_id) - try: - st = ts.get_structure(section_id) - except Exception: - return None - preview = str(st.get("preview") or "").replace("\n", " ").strip() - title = preview if preview else section_id - score = float(map_scores.get(section_id, 0.0) or 0.0) - return _MapNode( - section_id=section_id, - depth=depth, - title=title, - score=score, - n_lines=int(st.get("n_lines") or 0), - has_children=False, - parent_id=parent_id, - ) - - def append_visible_descendants( - parent_children: List[_MapNode], - section_id: str, - depth: int, - parent_id: Optional[str], - row_title_hint: Optional[dict] = None, - ) -> None: - """Attach section_id if visible. collected/dismissed drop node + subtree.""" - if not section_id or section_id in gone: - return - if keep is not None and section_id not in keep: - return - node = make_node(section_id, depth, parent_id) - if node is None: - return - if row_title_hint and (not node.title or node.title == section_id): - node.title = _title_from_row(row_title_hint, section_id) - if section_id in harvested: - # Collapsed leaf: this line alone represents the covered branch. - node.harvested_by = str(harvested[section_id]) - else: - for row in _children(ts, section_id): - child_id = str(row.get("section_id") or "").strip() - if child_id: - append_visible_descendants( - node.children, child_id, depth + 1, section_id, row - ) - node.has_children = bool(node.children) - _count_descendants(node) - parent_children.append(node) - - for rid in root_ids: - append_visible_descendants(roots, rid, 0, None) - return roots - - -def _clip_summary(text: str, *, head: int = 120) -> str: - s = (text or "").replace("\n", " ").strip() - if len(s) <= head: - return s - return s[: max(0, head - 1)].rstrip() + "…" - - -def format_hit_tag(*, is_highlight: bool) -> str: - """Render baseline Hit badge from highlight_ids.""" - return " [Hit]" if is_highlight else "" - - -def format_harvested_tag(owner_subgoal_id: str) -> str: - """Render the collapsed-coverage badge for a node kept visible post-collect.""" - sid = str(owner_subgoal_id or "").strip() - return f" [harvested:{sid}]" if sid else "" - - -def _render_map( - nodes: List[_MapNode], - *, - ts: Any = None, - doc_id: str, - scope: Optional[str], - char_limit: int, - highlight_ids: Optional[Set[str]] = None, - inline_summary: bool = False, -) -> tuple[str, List[SectionView], Dict[str, str], bool]: - """Render the budget-hidden title map (no mid-tree hard truncation).""" - _ = char_limit - hits = set(highlight_ids or ()) - lines: List[str] = [] - visible: List[SectionView] = [] - id_to_section: Dict[str, str] = {} - counter = 1 - any_hidden = False - - lines.append(f"doc_id={doc_id}") - lines.append(f"scope={scope or ''}") - lines.append( - "map=title+summary" if inline_summary else "map=title-only (action IDs attached per node)" - ) - - def _summary_doc_for(section_id: str) -> str: - if ts is None: - return str(doc_id or "") - try: - from .nav_address import owner_document - - return owner_document(ts, section_id, doc_id) or str(doc_id or "") - except Exception: - return str(doc_id or "") - - def render(node: _MapNode) -> None: - nonlocal counter, any_hidden - if node.hidden: - any_hidden = True - return - map_id = f"N{counter}" - counter += 1 - node.map_id = map_id - id_to_section[map_id] = node.section_id - indent = " " * node.depth - is_hit = node.section_id in hits or node.is_highlight - leaf_tag = " [Leaf]" if not node.has_children else "" - hit_tag = format_hit_tag(is_highlight=is_hit) - harvested_tag = format_harvested_tag(node.harvested_by) - line = f"{indent}[{map_id}] {node.title}{leaf_tag}{hit_tag}{harvested_tag}" - lines.append(line) - summary = "" - if inline_summary: - summary = _clip_summary( - _section_summary_for_map( - ts, - node.section_id, - doc_id=_summary_doc_for(node.section_id), - ) - ) - if summary: - lines.append(f"{indent} summary: {summary}") - visible.append( - SectionView( - section_id=node.section_id, - level=node.depth, - preview="", - score=node.score, - n_lines=node.n_lines, - has_children=node.has_children, - depth_from_scope=node.depth, - map_id=map_id, - title=node.title, - n_descendants=node.n_descendants, - is_highlight=is_hit, - parent_id=node.parent_id, - summary=summary, - harvested_by=node.harvested_by, - ) - ) - for child in node.children: - render(child) - - for root in nodes: - render(root) - - return "\n".join(lines), visible, id_to_section, any_hidden - - -def _expand_allowed_with_ancestors( - ts: Any, - allowed_section_ids: Optional[Set[str]], -) -> Optional[Set[str]]: - if allowed_section_ids is None: - return None - keep = {str(sid).strip() for sid in allowed_section_ids if str(sid).strip()} - parent_fn = getattr(ts, "parent_id", None) - if not callable(parent_fn): - return keep - extra: Set[str] = set() - for sid in keep: - cur = sid - seen: Set[str] = set() - while cur and cur not in seen: - seen.add(cur) - parent = parent_fn(cur) - if not parent: - break - extra.add(str(parent)) - cur = str(parent) - return keep | extra - - -def _fallback_highlights_from_tree(roots: List[_MapNode], k: int) -> List[str]: - leaves = [n for n in _flatten(roots, include_hidden=True) if not n.has_children] - leaves.sort(key=lambda n: (-n.score, n.section_id)) - return [n.section_id for n in leaves[: max(0, int(k))]] - - -def build_map( - ts: Any, - *, - doc_id: str, - query: str, - scope: Optional[str], - config: NavConfig, - map_scores: Optional[Dict[str, float]] = None, - collected_section_ids: Optional[Set[str]] = None, - dismissed_section_ids: Optional[Set[str]] = None, - highlight_ids: Optional[List[str]] = None, - extra_hidden_ids: Optional[Set[str]] = None, - harvested_section_ids: Optional[Dict[str, str]] = None, - allowed_section_ids: Optional[Set[str]] = None, -) -> Projection: - """Full-depth title map with score-ordered budget hiding (+ optional inline summary).""" - scores = dict(map_scores or {}) - if scope: - root_ids = [scope] - else: - root_ids = _top_sections(ts, doc_id) - keep_ids = _expand_allowed_with_ancestors(ts, allowed_section_ids) - roots = _build_map_tree( - ts, - root_ids=root_ids, - map_scores=scores, - collected_section_ids=collected_section_ids, - dismissed_section_ids=dismissed_section_ids, - harvested_section_ids=harvested_section_ids, - keep_ids=keep_ids, - ) - hits = [str(x).strip() for x in (highlight_ids or []) if str(x).strip()] - if not hits: - hits = _fallback_highlights_from_tree(roots, int(config.collect_top_k)) - hit_set = set(hits) - for node in _flatten(roots, include_hidden=True): - if node.section_id in hit_set: - node.is_highlight = True - - must_keep: Set[str] = set(hit_set) - for hid in hits: - must_keep.update(_ancestors_in_tree(roots, hid)) - - # Root is always title-only. A scoped region keeps summaries only while its - # full actionable map is small; large scopes become title-only so the agent - # can see more branches and choose another DISPATCH. - scope_summary_limit = max( - 0, int(getattr(config, "scope_inline_summary_char_limit", 2000) or 0) - ) - inline_summary = ( - scope is not None - and _estimate_actionable_total(roots, with_summary=True) - <= scope_summary_limit - ) - char_limit = max(1, int(config.map_char_limit)) - _apply_budget_hide( - roots, - char_limit=char_limit, - must_keep=must_keep, - extra_hidden_ids=extra_hidden_ids, - with_summary=inline_summary, - ) - text, tree_visible, id_map, truncated = _render_map( - roots, - ts=ts, - doc_id=doc_id, - scope=scope, - char_limit=char_limit, - highlight_ids=hit_set, - inline_summary=inline_summary, - ) - visible_sorted = sorted( - tree_visible, - key=lambda v: (-v.score, v.depth_from_scope, v.section_id), - ) - return Projection( - doc_id=doc_id, - scope=scope, - text=text, - visible_sections=visible_sorted, - truncated=truncated, - id_to_section=id_map, - tree_sections=list(tree_visible), - highlight_ids=list(hits), - ) - - -def top_visible_sections(projection: Projection, *, limit: int) -> List[SectionView]: - return list(projection.visible_sections[: max(0, limit)]) diff --git a/deprecated/mapnav/nav/nav_scope_filter.py b/deprecated/mapnav/nav/nav_scope_filter.py deleted file mode 100644 index a381a9a80..000000000 --- a/deprecated/mapnav/nav/nav_scope_filter.py +++ /dev/null @@ -1,426 +0,0 @@ -"""Bounded WHERE self-correction loop (pre-harvest). - -The policy writes or revises a ``NodeFilter``. Apply is deterministic. The -loop stops at ``filter_max_rounds``, token exhaustion, explicit fallback, or -agent ``done``. The agent may override settle; otherwise hits → scoped harvest. -""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from typing import Any, Dict, List, Literal, Optional, Sequence - -from shared.services.retrieval.scoring.node_filter_predicates import field_predicate - -from .nav_node_filter import ( - FilterResult, - NodeFilter, - apply_node_filter, - node_filter, - render_submap_observation, -) -from .nav_types import NavConfig - -ScopeDecision = Literal["collect_all", "scoped_harvest", "fallback"] -_SCOPE_FILTER_PURPOSE = "nav_scope_filter_v1" -_DECISIONS = {"collect_all", "scoped_harvest", "fallback"} - - -@dataclass -class ScopeFilterOutcome: - decision: ScopeDecision - settled_section_ids: List[str] = field(default_factory=list) - settled_doc_ids: List[str] = field(default_factory=list) - rounds: int = 0 - last_result: Optional[FilterResult] = None - reason: str = "" - - -def parse_node_filter(obj: Dict[str, Any] | None) -> Optional[NodeFilter]: - if not isinstance(obj, dict): - return None - raw = obj.get("predicates") - if raw is None: - raw = obj.get("filter") - if not isinstance(raw, list): - return None - preds = [] - for item in raw: - if not isinstance(item, dict): - continue - terms = item.get("terms") or [] - if isinstance(terms, str): - terms = [terms] - if not isinstance(terms, list): - continue - try: - preds.append( - field_predicate( - str(item.get("field") or ""), - [str(t) for t in terms], - str(item.get("match") or "substring"), - ) - ) - except ValueError: - continue - if not preds: - return None - return node_filter(preds) - - -def run_scope_filter( - ts: Any, - config: NavConfig, - *, - query: str, - doc_ids: Sequence[str], - map_observation: str = "", - seed_filter: Optional[NodeFilter] = None, - steps_out: Optional[List[Any]] = None, -) -> ScopeFilterOutcome: - """Apply → observe → revise until settle, fallback, or the round cap.""" - if not bool(getattr(config, "enable_node_filter", False)): - return ScopeFilterOutcome(decision="fallback", reason="disabled") - - from .nav_token_budget import nav_token_budget_exhausted, stamp_step_detail - - max_rounds = max(1, int(getattr(config, "filter_max_rounds", 3) or 3)) - wanted = [str(did).strip() for did in doc_ids if str(did).strip()] - map_text = str(map_observation or "").strip() or _compact_map(ts, wanted) - current = seed_filter - last_result: Optional[FilterResult] = None - last_obs = "" - last_decision: Optional[ScopeDecision] = None - last_reason = "" - - if current is None: - action = _scope_filter_policy_call( - config, - query=query, - full_map=map_text, - last_filter=None, - last_result=None, - last_observation="", - round_idx=0, - max_rounds=max_rounds, - include_full_map=True, - ) - current = action.get("filter") - last_decision = action.get("decision") - last_reason = str(action.get("reason") or "") - if action.get("kind") == "fallback" or current is None: - return ScopeFilterOutcome( - decision="fallback", - reason=last_reason or "policy_fallback", - ) - - for round_idx in range(1, max_rounds + 1): - if nav_token_budget_exhausted(): - return ScopeFilterOutcome( - decision="fallback", - settled_section_ids=list(last_result.matched_section_ids) - if last_result - else [], - settled_doc_ids=list(last_result.matched_doc_ids) if last_result else [], - rounds=round_idx - 1, - last_result=last_result, - reason="token_limit", - ) - assert current is not None - result = apply_node_filter(ts, wanted, current) - last_result = result - last_obs = render_submap_observation(ts, result, doc_ids=wanted) - if steps_out is not None: - from ._compat import AgentStep - - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="node_filter", - detail=stamp_step_detail( - { - "round": round_idx, - "predicates": _filter_payload(current), - "fields": sorted( - {p.field for p in current.predicates} - ), - "cardinality": result.cardinality, - "failed_predicates": list(result.failed_predicates), - "matched_section_ids": list(result.matched_section_ids), - "action": "", - "decision": "", - "reason": "", - } - ), - ) - ) - - is_last = round_idx >= max_rounds - if is_last: - if steps_out: - steps_out[-1].detail["action"] = "max_rounds" - return _settle( - result, - agent_decision=last_decision, - rounds=round_idx, - reason="max_rounds", - steps_out=steps_out, - ) - - action = _scope_filter_policy_call( - config, - query=query, - full_map=map_text, - last_filter=current, - last_result=result, - last_observation=last_obs, - round_idx=round_idx, - max_rounds=max_rounds, - include_full_map=False, - ) - if action.get("kind") == "widen": - # Agent judged the sub-map too narrow: one re-look with the full - # map, same round. A second widen settles on whatever it returns. - action = _scope_filter_policy_call( - config, - query=query, - full_map=map_text, - last_filter=current, - last_result=result, - last_observation=last_obs, - round_idx=round_idx, - max_rounds=max_rounds, - include_full_map=True, - ) - if action.get("kind") == "widen": - if action.get("filter") is not None: - action["kind"] = "filter" - else: - action["kind"] = "fallback" - action["reason"] = action.get("reason") or "widen_without_filter" - last_decision = action.get("decision") - last_reason = str(action.get("reason") or "") - kind = str(action.get("kind") or "") - if steps_out: - steps_out[-1].detail["action"] = kind - steps_out[-1].detail["reason"] = last_reason - if kind == "fallback": - if steps_out: - steps_out[-1].detail["decision"] = "fallback" - return ScopeFilterOutcome( - decision="fallback", - settled_section_ids=list(result.matched_section_ids), - settled_doc_ids=list(result.matched_doc_ids), - rounds=round_idx, - last_result=result, - reason=last_reason or "policy_fallback", - ) - if kind == "done": - return _settle( - result, - agent_decision=last_decision, - rounds=round_idx, - reason=last_reason or "done", - steps_out=steps_out, - ) - nxt = action.get("filter") - if nxt is not None: - current = nxt - - assert last_result is not None - return _settle( - last_result, - agent_decision=last_decision, - rounds=max_rounds, - reason=last_reason or "max_rounds", - steps_out=steps_out, - ) - - -def _settle( - result: FilterResult, - *, - agent_decision: Optional[ScopeDecision], - rounds: int, - reason: str, - steps_out: Optional[List[Any]], -) -> ScopeFilterOutcome: - if result.cardinality <= 0: - decision: ScopeDecision = "fallback" - settle_reason = reason or "no_hits" - elif agent_decision in _DECISIONS: - decision = agent_decision - settle_reason = reason or "agent" - if decision == "fallback": - settle_reason = reason or "agent_fallback" - else: - decision = "scoped_harvest" - settle_reason = reason or "hits" - if steps_out: - steps_out[-1].detail["decision"] = decision - steps_out[-1].detail["reason"] = settle_reason - return ScopeFilterOutcome( - decision=decision, - settled_section_ids=list(result.matched_section_ids), - settled_doc_ids=list(result.matched_doc_ids), - rounds=rounds, - last_result=result, - reason=settle_reason, - ) - - -def _scope_filter_policy_call( - config: NavConfig, - *, - query: str, - full_map: str, - last_filter: Optional[NodeFilter], - last_result: Optional[FilterResult], - last_observation: str, - round_idx: int, - max_rounds: int, - include_full_map: bool, -) -> Dict[str, Any]: - from .nav_llm import nav_chat, resolve_nav_model - from .nav_policy import _extract_json_obj - from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted - - if nav_token_budget_exhausted(): - return {"kind": "fallback", "reason": "token_limit"} - - model = resolve_nav_model( - model=config.llm_model, - model_env="NAV_LLM_MODEL", - fallback_envs=("NAV_LLM_MODEL",), - ) - user = f"Query: {query}\nRound: {round_idx}/{max_rounds}\n" - if last_filter is not None: - user += f"Last filter: {json.dumps(_filter_payload(last_filter), ensure_ascii=False)}\n" - if last_result is not None: - user += f"Last hits: {last_result.cardinality}\n" - if last_observation: - user += f"=== Sub-map (hits of last filter) ===\n{last_observation}\n" - if include_full_map: - user += f"=== Full map ===\n{full_map}\n=== End Full Map ===\n" - try: - cached = nav_chat( - purpose=_SCOPE_FILTER_PURPOSE, - model=model, - messages=[ - {"role": "system", "content": _scope_filter_system_prompt(seed=last_filter is None)}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=max(256, int(config.llm_max_tokens or 256)), - response_format={"type": "json_object"}, - context="Nav Scope Filter", - usage_tag="nav_scope_filter", - ) - except NavTokenLimit: - return {"kind": "fallback", "reason": "token_limit"} - - text = str(cached.get("content") or "").strip() - obj = _extract_json_obj(text) or {} - kind = str(obj.get("action") or obj.get("kind") or "filter").strip().lower() - if kind not in {"filter", "done", "widen", "fallback"}: - kind = "filter" - decision_raw = str(obj.get("decision") or "").strip().lower() - decision: Optional[ScopeDecision] = ( - decision_raw if decision_raw in _DECISIONS else None - ) - parsed = parse_node_filter(obj) - return { - "kind": kind, - "filter": parsed, - "decision": decision, - "reason": str(obj.get("reason") or "")[:300], - "raw": text, - } - - -def _scope_filter_system_prompt(*, seed: bool) -> str: - field_schema = ( - '"path"' if seed else '"path|summary"' - ) - field_rule = ( - "Match on section paths only (filenames and title chains)." - if seed - else "Fields AND together; terms inside one field OR together." - ) - base = ( - "You write a WHERE filter over document sections. Return json.\n" - "Schema: {\"action\":\"filter|done|widen|fallback\",\"predicates\":" - f"[{{\"field\":{field_schema},\"terms\":[\"...\"],\"match\":\"substring|regex\"}}]," - "\"decision\":\"collect_all|scoped_harvest|fallback\",\"reason\":\"...\"}\n" - f"{field_rule}\n" - "Write terms as natural-language words from the query or the map " - "(entities, topics, aliases); do not rely on section numbering alone.\n" - ) - if seed: - return base + ( - "Write the first filter for the query against the full map. " - "action=filter returns it; action=fallback only when no path " - "predicate can isolate the target sections." - ) - return base + ( - "You see the sub-map hit by the last filter; judge by its content:\n" - "- sub-map covers the query -> action=done\n" - "- sub-map has off-topic nodes -> action=filter with a narrower " - "revision of the last filter\n" - "- sub-map looks too narrow or misses parts of the query -> " - "action=widen to see the full map once, then revise\n" - "Revise the last filter, never restart from the query. " - "action=fallback only when no predicate can isolate the target sections." - ) - - -def _filter_payload(nf: NodeFilter) -> List[Dict[str, Any]]: - return [ - {"field": p.field, "terms": list(p.terms), "match": p.match} - for p in nf.predicates - ] - - -def _compact_map(ts: Any, doc_ids: Sequence[str]) -> str: - path_fn = getattr(ts, "path_titles", None) - structure_fn = getattr(ts, "get_structure", None) - lines: List[str] = [] - - def path_of(sid: str, doc_id: str) -> str: - if not callable(path_fn): - return sid - try: - return str(path_fn(sid, doc_id) or sid) - except TypeError: - return str(path_fn(sid) or sid) - - def summary_of(sid: str) -> str: - if not callable(structure_fn): - return "" - try: - st = structure_fn(sid) or {} - return str(st.get("summary") or "").strip() - except Exception: - return "" - - provider = getattr(ts, "_provider", None) - child_fn = getattr(provider, "children", None) if provider is not None else None - root_fn = getattr(ts, "sections_for_doc", None) - - for doc_id in doc_ids: - lines.append(path_of(doc_id, doc_id)) - stack = [str(s) for s in (root_fn(doc_id) if callable(root_fn) else []) if str(s)] - seen: set[str] = set() - while stack: - sid = stack.pop(0) - if not sid or sid in seen: - continue - seen.add(sid) - path = path_of(sid, doc_id) - summary = summary_of(sid) - line = path if not summary else f"{path} | {summary}" - lines.append(line) - kids = [str(c) for c in (child_fn(sid) if callable(child_fn) else []) if str(c)] - stack[0:0] = kids - return "\n".join(lines) diff --git a/deprecated/mapnav/nav/nav_token_budget.py b/deprecated/mapnav/nav/nav_token_budget.py deleted file mode 100644 index afae352df..000000000 --- a/deprecated/mapnav/nav/nav_token_budget.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Episode LLM token hard stop. - -One counter, one rule: if cumulative LLM ``total_tokens`` reaches the episode -``token_limit`` (from ``NavConfig``, else ``RETRIEVAL_NAV_TOKEN_LIMIT`` / default), -do not start another nav LLM call; keep already-collected evidence and return -upward (FINISH / done / fallback). - -Inside ``nav_token_episode()`` the counter is per-episode (contextvars). -Outside it, usage falls back to the process-global ``llm_usage`` snapshot -so bin scripts that only ``reset_usage()`` keep working. -""" - -from __future__ import annotations - -import os -import time -from contextlib import contextmanager -from contextvars import ContextVar -from typing import Any, Dict, Iterator, List, Optional - -_ENV_TOKEN_LIMIT = "RETRIEVAL_NAV_TOKEN_LIMIT" -_DEFAULT_TOKEN_LIMIT = 100_000 - -# Mutable one-element list so callers can increment without rebinding ContextVar. -_EpisodeCounter = List[int] -_episode_tokens: ContextVar[Optional[_EpisodeCounter]] = ContextVar( - "nav_episode_tokens", default=None -) -_episode_limit: ContextVar[Optional[int]] = ContextVar( - "nav_episode_token_limit", default=None -) -_last_stamp_tokens: ContextVar[Optional[int]] = ContextVar( - "nav_last_stamp_tokens", default=None -) - - -class NavTokenLimit(Exception): - """Raised when the episode (or process) LLM token budget is exhausted.""" - - def __init__(self, used: int = 0, limit: int = 0) -> None: - self.used = int(used) - self.limit = int(limit) - super().__init__( - f"nav token limit exhausted: used={self.used} limit={self.limit}" - ) - - -def nav_token_limit() -> int: - """Always a positive limit: episode bind → env → default.""" - bound = _episode_limit.get() - if bound is not None and int(bound) > 0: - return int(bound) - try: - limit = int(os.environ.get(_ENV_TOKEN_LIMIT, "").strip()) - except ValueError: - limit = 0 - return limit if limit > 0 else _DEFAULT_TOKEN_LIMIT - - -def _process_tokens_used() -> int: - from ._compat import snapshot_usage # type: ignore - - total = 0 - for block in snapshot_usage().values(): - total += int(block.get("total_tokens", 0) or 0) - return total - - -def nav_tokens_used() -> int: - ep = _episode_tokens.get() - if ep is not None: - return int(ep[0]) - return _process_tokens_used() - - -def nav_token_budget_exhausted() -> bool: - return nav_tokens_used() >= nav_token_limit() - - -def record_episode_tokens(usage: Optional[Dict[str, Any]]) -> None: - """Add one call's ``total_tokens`` to the active episode counter (no-op outside).""" - ep = _episode_tokens.get() - if ep is None: - return - try: - add = int((usage or {}).get("total_tokens", 0) or 0) - except (TypeError, ValueError): - add = 0 - if add > 0: - ep[0] += add - - -def stamp_step_detail( - detail: Optional[Dict[str, Any]] = None, - *, - t0: Optional[float] = None, -) -> Dict[str, Any]: - """Annotate an AgentStep detail with episode token counters and step elapsed_ms.""" - stamped = dict(detail or {}) - used = int(nav_tokens_used()) - prev = _last_stamp_tokens.get() - if prev is None: - prev = 0 - stamped["token_limit"] = int(nav_token_limit()) - stamped["tokens_used_total"] = used - stamped["tokens_used_delta"] = max(0, used - int(prev)) - if t0 is not None: - stamped["elapsed_ms"] = int(round((time.perf_counter() - float(t0)) * 1000.0)) - _last_stamp_tokens.set(used) - return stamped - - -@contextmanager -def nav_token_episode(*, token_limit: int = 0) -> Iterator[None]: - """Bind a fresh per-episode token counter (and optional hard limit).""" - token = _episode_tokens.set([0]) - stamp_token = _last_stamp_tokens.set(0) - limit_arg = int(token_limit or 0) - limit_token = _episode_limit.set(limit_arg if limit_arg > 0 else None) - try: - yield - finally: - _episode_limit.reset(limit_token) - _last_stamp_tokens.reset(stamp_token) - _episode_tokens.reset(token) diff --git a/deprecated/mapnav/nav/nav_types.py b/deprecated/mapnav/nav/nav_types.py deleted file mode 100644 index 323dac6a9..000000000 --- a/deprecated/mapnav/nav/nav_types.py +++ /dev/null @@ -1,281 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum -from typing import Any, Dict, List, Literal, Optional, Tuple - -NavMode = Literal["checklist"] - - -class ActionKind(str, Enum): - COLLECT = "collect" - DISPATCH = "dispatch" - FINISH = "finish" - - -@dataclass -class NavConfig: - collect_top_k: int = 6 # rescue-K for highlights (not action quota) - read_score_bonus: float = 10.0 - policy: str = "rule" - # Empty → resolve_nav_model falls back to NAV_* env (EXP scripts). - # Knowhere build_nav_config sets concrete model names. - llm_model: str = "" - llm_temperature: float = 0.0 - llm_max_tokens: int = 256 - map_char_limit: int = 5000 # display budget (fold threshold); only hard display limit - # Recursive dispatch. - enable_recursive_dispatch: bool = True - max_dispatch_depth: int = 5 - subagent_model: str = "" - # Scoped maps whose estimated (with-summary) size exceeds this threshold drop - # inline summaries (title-only), nudging the agent to DISPATCH deeper rather - # than broadly COLLECT the whole parent. run_nav_episode re-derives it from - # the episode's evidence budget (budget_chars x _SCOPE_SUMMARY_BUDGET_MULT). - scope_inline_summary_char_limit: int = 1500 - # COMPOSE child score = own_unit + compose_confidence_weight * collect_confidence - # (see nav_compose._child_final_score); drives group_key / within-group rank. - compose_confidence_weight: float = 0.5 - # Product mode: checklist = plan+harvest+control. - mode: NavMode = "checklist" - # Display budget for the one-shot planning map; executor still uses map_char_limit. - # 0 = reuse map_char_limit. - planning_map_char_limit: int = 10000 - # Soft prompt guidance only when > 0; never silently truncates a valid plan. - planner_max_subgoals: int = 0 - planner_model: str = "" - # "", "enabled", or "disabled". Empty → NAV_PLANNER_THINKING / default off. - planner_thinking: str = "" - planner_think_max_tokens: int = 0 - # 0 → RETRIEVAL_NAV_TOKEN_LIMIT env / default 100000. - token_limit: int = 0 - # Plan JSON is larger than a single harvest action; give the planner more room. - planner_llm_max_tokens: int = 1024 - # Harvest multi-id JSON (collect_ids + per-id confidence) needs more than a - # bare 256 — capped completion truncates mid-object and parses as empty. - harvest_llm_max_tokens: int = 1024 - # Checklist: harvest cycles per subgoal before drop. Min 1. - subgoal_max_attempts: int = 2 - # Checklist: 0 = never replan; otherwise hard cap on structural replans. - # Default 1 so shared-checklist probe/harness match without silent overrides. - max_replans: int = 1 - # Checklist: 0 = no extra wave cap (stop when no ready subgoals). - max_waves: int = 0 - # Structural recursion depth cap for harvest() (checklist mode). - max_harvest_depth: int = 5 - # WHERE node filter (pre-harvest). Off until orchestrate enables a subgoal. - enable_node_filter: bool = False - filter_max_rounds: int = 3 - - @property - def is_checklist(self) -> bool: - return str(self.mode or "").strip().lower() == "checklist" - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> "NavConfig": - flat = dict(data) - # Retired product flags (dropped after the navigate loop was removed). - for dead in ( - "expand_top_k", - "map_peek_top_k", - "map_jump_top_k", - "peek_content_fanout", - "peek_content_chars", - "map_collapse_min_score", - "dispatch_group_size", - "dispatch_max_workers", - "projection_depth", - "projection_child_limit", - "projection_char_limit", - "enable_contract_verify", - "enable_per_subgoal_illumination", - "enable_goal_conditioned_folding", - "enable_subgoal_budget_ledger", - "subgoal_budget_floor_frac", - "enable_anchor_entry", - "enable_settle_group_rank", - "enable_external_rerank", - "compose_preview_snippet_chars", - "compose_preview_max_children", - "compose_packing_mode", - "compose_coverage_budget_frac", - "compose_snippet_chars", - "enable_query_planning", - "enable_plan_orchestration", - "enable_slot_extract", - "enable_one_shot_harvest", - "enable_plan_control", - "show_harvested_in_map", - "dispatch_concurrency", - # Retired: config held env *names*; now holds model/thinking values. - "llm_model_env", - "planner_model_env", - "subagent_model_env", - # Retired with the multi-step navigate loop. - "navigate_max_steps", - "enable_depth0_oversize_to_dispatch", - "depth0_oversize_char_limit", - "compose_group_rank_max_chars", - # Never wired to any consumer. - "search_k", - "plan_control_digest_chars", - # Legacy shallow-projection leftovers. - "summary_chars", - "map_children_limit", - # Budget-mode step gating: build_legal_actions always ran at - # step_idx=0, so critical/tight never triggered. - "max_steps", - "budget_modes", - "critical_remaining_steps", - "tight_remaining_steps", - # Folded into nav_agent._SCOPE_SUMMARY_BUDGET_MULT. - "scope_inline_summary_budget_mult", - "collect_k", - "filter_min_hits", - "filter_max_hits", - "map_mode", - ): - flat.pop(dead, None) - flat["mode"] = "checklist" - allowed = {f.name for f in cls.__dataclass_fields__.values()} - cfg = cls(**{k: v for k, v in flat.items() if k in allowed}) - if cfg.llm_max_tokens < 256: - cfg.llm_max_tokens = 256 - return cfg - - -@dataclass -class SectionView: - section_id: str - level: int - preview: str - score: float = 0.0 - n_lines: int = 0 - has_children: bool = False - depth_from_scope: int = 0 - map_id: str = "" - title: str = "" - n_descendants: int = 0 - is_highlight: bool = False - parent_id: Optional[str] = None - summary: str = "" - # PLAN×NAV: subgoal id that collected this node's branch, when the - # node stays visible (collapsed, no descendant expansion) instead of being - # deleted from the map — for [harvested:sN]. - harvested_by: str = "" - - -@dataclass -class Projection: - doc_id: str - scope: Optional[str] - text: str - visible_sections: List[SectionView] - truncated: bool = False # True if any budget-hidden nodes - id_to_section: Dict[str, str] = field(default_factory=dict) - tree_sections: List[SectionView] = field(default_factory=list) - highlight_ids: List[str] = field(default_factory=list) - - -@dataclass -class LegalAction: - action_id: str - kind: ActionKind - section_id: Optional[str] = None - query: str = "" - label: str = "" - score: float = 0.0 - metadata: Dict[str, Any] = field(default_factory=dict) - - def prompt_line(self) -> str: - target = self.section_id or "" - bits = [self.action_id, self.kind.value.upper()] - if target: - bits.append(target) - if self.label: - bits.append(self.label) - if self.score: - bits.append(f"score={self.score:.4f}") - return " | ".join(bits) - - -@dataclass -class SubgoalResult: - """Typed outcome of one subgoal execution (M5).""" - - subgoal_id: str - satisfied: bool - confidence: float = 0.0 - collected_section_ids: List[str] = field(default_factory=list) - # Explicit COLLECT targets this wave only (hydration descendants omitted). - explicit_collect_ids: List[str] = field(default_factory=list) - extracted: Dict[str, str] = field(default_factory=dict) - gap: str = "" - chars_used: int = 0 - - -@dataclass -class NavState: - doc_id: str - query: str - task_type: str = "unknown" - # Working scope for the current harvest level (set by harvest(), not a stack). - current_scope: Optional[str] = None - collected_ids: set[str] = field(default_factory=set) - collected: List[Tuple[Any, float]] = field(default_factory=list) - map_scores: Dict[str, float] = field(default_factory=dict) - unit_scores: Dict[str, float] = field(default_factory=dict) - highlight_ids: List[str] = field(default_factory=list) - # "Branch done / removed from map" = COLLECT'd sid ∪ all descendants. - collected_section_ids: set[str] = field(default_factory=set) - blocked_collect_section_ids: set[str] = field(default_factory=set) - action_history: List[Dict[str, Any]] = field(default_factory=list) - dismissed_section_ids: set[str] = field(default_factory=set) - # Explicit COLLECT confidence by section_id; hydration-only descendants stay 0. - collect_confidence: Dict[str, float] = field(default_factory=dict) - # Explicit COLLECT targets only (batch action sids); hydration descendants omitted. - # Used by plan_control digest and compose progressive trim (protect explicit owners). - explicit_collect_ids: set[str] = field(default_factory=set) - # External agent relative priority over nearest-parent groups: parent_id -> score - # (higher packs first). Empty = no external rerank yet. - group_priority: Dict[str, float] = field(default_factory=dict) - # M2: structured retrieval plan + slot bindings for delayed query fill. - retrieval_plan: Optional[Any] = None - slot_bindings: Dict[str, str] = field(default_factory=dict) - satisfied_subgoal_ids: set[str] = field(default_factory=set) - # Finished trying (accept or drop). Widen leaves a subgoal out of this set. - attempted_subgoal_ids: set[str] = field(default_factory=set) - # Soft focus for policy (never clips action space). - focus_subgoal_id: str = "" - focus_subgoal_need: str = "" - focus_subgoal_contract: str = "" - focus_retrieval_query: str = "" - focus_contract_kind: str = "" - subgoal_results: Dict[str, Any] = field(default_factory=dict) - replan_count: int = 0 - # Checklist: explicit collect-root section_id -> owning subgoal_id - # (drives "[harvested:sN]" map tags). - harvested_owner_subgoal: Dict[str, str] = field(default_factory=dict) - # Last widen gap note per subgoal (what plan_control said was missing). - # Input to the PLAN-side query rewrite; never concatenated into a query. - subgoal_widen_gaps: Dict[str, str] = field(default_factory=dict) - # PLAN's rewritten retrieval_query per subgoal after widen. Overrides the - # planned query for the next harvest, and re-scores the shared map with it. - subgoal_refined_queries: Dict[str, str] = field(default_factory=dict) - # Episode-local map scores keyed by retrieval query. Checklist waves may - # revisit the same subgoal query; reuse the exact score snapshot instead of - # rebuilding the persisted index and rescoring the corpus. - relit_map_cache: Dict[ - str, Tuple[Dict[str, float], Dict[str, float], List[str]] - ] = field(default_factory=dict) - # Per-subgoal "seen but not selected" section ids — hidden from later map - # views for that subgoal so widen surfaces siblings instead of dead ends. - subgoal_dismissed_section_ids: Dict[str, set[str]] = field(default_factory=dict) - # Per-subgoal wave-attempt counter (circuit breaker under plan_control). - subgoal_attempt_counts: Dict[str, int] = field(default_factory=dict) - # Terminal "drop" outcomes: disjoint from satisfied_subgoal_ids. Union of - # the two is "settled" for dependency readiness. - dropped_subgoal_ids: set[str] = field(default_factory=set) - # SEARCH_* inspector results injected into the next harvest observation - # (Knowhere asset tool context). Cleared/overwritten by apply_search_assets. - asset_observation_context: str = "" diff --git a/deprecated/mapnav/nav/nav_verify.py b/deprecated/mapnav/nav/nav_verify.py deleted file mode 100644 index aa198ab96..000000000 --- a/deprecated/mapnav/nav/nav_verify.py +++ /dev/null @@ -1,259 +0,0 @@ -"""Slot binding helpers for plan orchestration. - -Extracts values only when a later subgoal's ``{{sN.slot}}`` (or short -``{{slot}}`` matching this subgoal's ``produces``) needs them. -Checklist reconciliation belongs solely to ``nav_control.plan_control``. -Slot fill is LLM-only (no lexical/heuristic line dump). -""" - -from __future__ import annotations - -import json -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple - -from .nav_plan import RetrievalPlan, Subgoal, unbound_slots -from .nav_types import NavConfig, SubgoalResult - - -def evidence_chars(text: str) -> int: - return len((text or "").strip()) - - -def _slot_owner_and_name(ref: str) -> Tuple[Optional[str], str]: - token = (ref or "").strip() - if not token: - return None, "" - if "." in token: - owner, name = token.split(".", 1) - return owner.strip() or None, name.strip() - return None, token - - -def demanded_slot_names(plan: RetrievalPlan, producer: Subgoal) -> List[str]: - """Slots this subgoal must fill because a *later* subgoal references them.""" - wanted: Set[str] = set() - produces = {str(p).strip() for p in (producer.produces or []) if str(p).strip()} - for sg in plan.subgoals: - if sg.id == producer.id: - continue - for ref in unbound_slots(sg.retrieval_query) + unbound_slots(sg.need): - owner, name = _slot_owner_and_name(ref) - if not name: - continue - if owner == producer.id: - wanted.add(name) - elif owner is None and name in produces: - wanted.add(name) - if produces: - # Preserve planner order; drop produces nobody consumes. - return [p for p in (producer.produces or []) if str(p).strip() in wanted] - return sorted(wanted) - - -def extract_slots_llm( - slots: Sequence[str], - evidence_text: str, - config: NavConfig, - *, - need: str = "", - retrieval_query: str = "", - contract_kind: str = "single_fact", - cardinality: Optional[int] = None, - steps_out: Optional[List[Any]] = None, - subgoal_id: str = "", -) -> Dict[str, str]: - """Ask the LLM for slot values; empty on failure.""" - names = [str(s).strip() for s in slots if str(s).strip()] - if not names or not (evidence_text or "").strip(): - return {} - try: - from .nav_llm import nav_chat, resolve_nav_model - from .nav_token_budget import nav_token_budget_exhausted, stamp_step_detail - except Exception: - return {} - if nav_token_budget_exhausted(): - return {} - - model = resolve_nav_model( - model=config.planner_model, - model_env="NAV_PLANNER_MODEL", - fallback_envs=("NAV_LLM_MODEL",), - ) - system = ( - "Extract slot values for a retrieval subgoal from evidence text.\n" - 'Return ONLY JSON: {"slots": {"name": "value"}, "confidence": 0..1}.\n' - "Use the evidence language. If a slot is missing, omit it.\n" - "Do not invent facts absent from the evidence." - ) - user = ( - f"Need: {need}\n" - f"Retrieval query: {retrieval_query or need}\n" - f"Slots to fill: {json.dumps(names, ensure_ascii=False)}\n" - f"Contract: {contract_kind}" - + (f", cardinality={cardinality}" if cardinality is not None else "") - + f"\n\n=== Evidence ===\n{evidence_text}\n=== End Evidence ===\n" - ) - import time - - t0 = time.perf_counter() - out: Dict[str, str] = {} - try: - cached = nav_chat( - purpose="nav_slot_extract_v1", - model=model, - messages=[ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=max(256, int(config.llm_max_tokens or 256)), - response_format={"type": "json_object"}, - context="Nav Slot Extract", - api_key_env="NAV_PLANNER_API_KEY", - base_url_env="NAV_PLANNER_BASE_URL", - usage_tag="nav_slot_extract", - ) - raw = str(cached.get("content") or "").strip() - obj = json.loads(raw) if raw.startswith("{") else {} - slots_obj = obj.get("slots") if isinstance(obj, dict) else None - if isinstance(slots_obj, dict): - for name in names: - val = str(slots_obj.get(name) or "").strip() - if val: - out[name] = val - except Exception: - out = {} - if steps_out is not None: - try: - from ._compat import AgentStep # type: ignore - except Exception: - AgentStep = None # type: ignore - if AgentStep is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="slot_extract", - detail=stamp_step_detail( - { - "subgoal_id": subgoal_id, - "slots_requested": names, - "slots_filled": sorted(out.keys()), - }, - t0=t0, - ), - ) - ) - return out - - -def extract_slots( - plan: RetrievalPlan, - subgoal: Subgoal, - evidence_text: str, - config: NavConfig, - *, - retrieval_query: str = "", - steps_out: Optional[List[Any]] = None, -) -> Tuple[Dict[str, str], float]: - """LLM-fill slots demanded by downstream subgoals; empty if none demanded.""" - slots = demanded_slot_names(plan, subgoal) - if not slots: - conf = 1.0 if (evidence_text or "").strip() else 0.0 - return {}, conf - filled = extract_slots_llm( - slots, - evidence_text, - config, - need=subgoal.need, - retrieval_query=retrieval_query or subgoal.retrieval_query, - contract_kind=subgoal.contract.kind, - cardinality=subgoal.contract.cardinality, - steps_out=steps_out, - subgoal_id=str(subgoal.id or ""), - ) - n = sum(1 for s in slots if (filled.get(s) or "").strip()) - return filled, float(n) / float(len(slots)) - - -def build_subgoal_result( - plan: RetrievalPlan, - state_collected_section_ids: Sequence[str], - config: NavConfig, - subgoal: Subgoal, - *, - retrieval_query: str, - new_chunks: Sequence[Tuple[Any, float]], - collected_before: Set[str], - explicit_collect_ids: Optional[Sequence[str]] = None, - explicit_before: Optional[Set[str]] = None, - steps_out: Optional[List[Any]] = None, -) -> SubgoalResult: - """Package this wave's evidence + optional downstream slot bindings. - - ``satisfied`` here means "this wave collected non-empty evidence" — bookkeeping - for dependency readiness when plan_control is off. Checklist acceptance is - decided only by ``plan_control``. - """ - evidence = build_evidence_text_from_chunks(new_chunks) - chars = evidence_chars(evidence) - extracted, conf = extract_slots( - plan, - subgoal, - evidence, - config, - retrieval_query=retrieval_query, - steps_out=steps_out, - ) - before_explicit = set(explicit_before or ()) - explicit_wave = [ - str(s).strip() - for s in (explicit_collect_ids or ()) - if str(s).strip() and str(s).strip() not in before_explicit - ] - seen: Set[str] = set() - explicit_ordered: List[str] = [] - for sid in explicit_wave: - if sid in seen: - continue - seen.add(sid) - explicit_ordered.append(sid) - return SubgoalResult( - subgoal_id=subgoal.id, - satisfied=chars > 0, - confidence=float(conf), - extracted=dict(extracted), - chars_used=chars, - gap="empty_evidence" if chars <= 0 else "", - collected_section_ids=[ - s for s in state_collected_section_ids if s not in collected_before - ], - explicit_collect_ids=explicit_ordered, - ) - - -def apply_bindings_from_result( - bindings: Dict[str, str], - subgoal: Subgoal, - extracted: Dict[str, str], -) -> Dict[str, str]: - """Write short and qualified ``s1.slot`` keys into bindings.""" - out = dict(bindings or {}) - for slot, value in (extracted or {}).items(): - name = str(slot).strip() - val = str(value or "").strip() - if not name or not val: - continue - out[name] = val - out[f"{subgoal.id}.{name}"] = val - return out - - -def build_evidence_text_from_chunks(chunks: Any) -> str: - """Concatenate (chunk, score) texts in order.""" - parts: List[str] = [] - for chunk, _score in list(chunks or []): - text = str(getattr(chunk, "text", "") or getattr(chunk, "content", "") or "") - if not text.strip(): - continue - parts.append(text) - return "\n".join(parts) diff --git a/deprecated/mapnav/nav_bridge.py b/deprecated/mapnav/nav_bridge.py deleted file mode 100644 index 677f40f63..000000000 --- a/deprecated/mapnav/nav_bridge.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Map-nav exit bridge: kept_chunks → referenced_chunks (+ scores). - -Refs always come from ``NavSnapshot.chunk_ref_index`` (DB-original -section_path / job_id), never from remounted provider paths. -""" - -from __future__ import annotations - -from typing import Any, Iterable, Optional - -from shared.services.retrieval.nav_snapshot import NavSnapshot - - -def dedupe_referenced_chunks(refs: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: - """Inline of workflow/reference_projection.WorkflowReferenceProjection.dedupe.""" - seen: set[str] = set() - out: list[dict[str, Any]] = [] - for ref in refs: - document_id = str(ref.get("document_id") or "").strip() - chunk_id = str(ref.get("chunk_id") or "").strip() - section_path = str(ref.get("section_path") or "").strip() - file_path = str(ref.get("file_path") or "").strip() - key = ( - f"{document_id}:{chunk_id}:{section_path}:{file_path}" - if document_id and chunk_id - else str(ref) - ) - if key in seen: - continue - seen.add(key) - out.append(dict(ref)) - return out - - -def _lookup_chunk_meta( - snapshot: NavSnapshot, - chunk_id: str, - *, - document_id: Optional[str] = None, -) -> Optional[dict[str, Any]]: - doc = str(document_id or "").strip() - if doc: - scoped = snapshot.chunk_ref_index.get(f"{doc}:{chunk_id}") - if isinstance(scoped, dict): - return scoped - meta = snapshot.chunk_ref_index.get(chunk_id) - return meta if isinstance(meta, dict) else None - - -def _ref_for( - chunk_id: str, - *, - snapshot: NavSnapshot, - score: Optional[float], - document_id: Optional[str] = None, -) -> Optional[dict[str, Any]]: - meta = _lookup_chunk_meta(snapshot, chunk_id, document_id=document_id) - if not meta: - return None - ref: dict[str, Any] = { - "chunk_id": chunk_id, - "document_id": meta.get("document_id"), - "chunk_type": meta.get("chunk_type") or "text", - "section_path": meta.get("section_path"), - "file_path": meta.get("file_path") or "", - "job_id": meta.get("job_id"), - } - if score is not None: - ref["score"] = float(score) - return ref - - -def _expand_node_to_chunk_ids( - *, - node_id: str, - section_id: Optional[str], - snapshot: NavSnapshot, - document_id: Optional[str] = None, -) -> list[str]: - nid = str(node_id or "").strip() - if not nid: - return [] - if _lookup_chunk_meta(snapshot, nid, document_id=document_id) is not None: - return [nid] - - sid = str(section_id or "").strip() or nid - if sid.endswith("__self"): - sid = sid[: -len("__self")] - elif sid.endswith("__path"): - sid = sid[: -len("__path")] - if nid.endswith("__self"): - sid = nid[: -len("__self")] - elif nid.endswith("__path"): - sid = nid[: -len("__path")] - - out: list[str] = [] - for unit in snapshot.provider.self_units(sid) or (): - cid = str(getattr(unit, "chunk_id", "") or "").strip() - if cid and _lookup_chunk_meta(snapshot, cid, document_id=document_id) is not None: - out.append(cid) - return out - - -def build_referenced_chunks( - episode: Any, - snapshot: NavSnapshot, -) -> tuple[list[dict[str, Any]], dict[str, float]]: - """Expand ``episode.kept_chunks`` into API refs + score index. - - Returns ``(refs, score_by_chunk_id)``. - """ - score_by_node: dict[str, float] = {} - for item in getattr(episode, "scored_chunks", None) or (): - if not isinstance(item, (tuple, list)) or len(item) < 2: - continue - chunk, score = item[0], item[1] - nid = str(getattr(chunk, "node_id", "") or "").strip() - if not nid: - continue - try: - score_by_node[nid] = float(score) - except (TypeError, ValueError): - continue - - refs: list[dict[str, Any]] = [] - score_by_chunk_id: dict[str, float] = {} - for chunk in getattr(episode, "kept_chunks", None) or (): - nid = str(getattr(chunk, "node_id", "") or "").strip() - sid = getattr(chunk, "section_id", None) - doc_id = str(getattr(chunk, "doc_id", "") or "").strip() or None - node_score = score_by_node.get(nid) - for chunk_id in _expand_node_to_chunk_ids( - node_id=nid, - section_id=str(sid) if sid else None, - snapshot=snapshot, - document_id=doc_id, - ): - if node_score is not None: - prev = score_by_chunk_id.get(chunk_id) - score_by_chunk_id[chunk_id] = ( - max(prev, node_score) if prev is not None else node_score - ) - ref = _ref_for( - chunk_id, - snapshot=snapshot, - score=score_by_chunk_id.get(chunk_id), - document_id=doc_id, - ) - if ref is not None: - refs.append(ref) - - return dedupe_referenced_chunks(refs), score_by_chunk_id diff --git a/deprecated/mapnav/nav_config.py b/deprecated/mapnav/nav_config.py deleted file mode 100644 index c7338f095..000000000 --- a/deprecated/mapnav/nav_config.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Production map-nav config — single source of truth at the retrieval entry. - -Values live on ``NavConfig`` (model / thinking / token_limit / evidence chars). -Vendored nav binds them for the episode via ``nav_llm_runtime`` + -``nav_token_episode``; no process-wide env seeding and no Knowhere wallet/agentic knobs. - -LEGACY, PENDING REPLACEMENT: this config only feeds the map-nav -PLANNER/HARVEST/CONTROL episode -(``shared.services.retrieval.nav``). It is not read by ``agent_explore`` -and is scheduled for Phase 5 cleanup once map-nav is -no longer the default route. -""" - -from __future__ import annotations - -from typing import Any - -from shared.services.retrieval.nav.nav_types import NavConfig -from shared.services.retrieval.settings import EVIDENCE_TEXT_CHAR_BUDGET - -# Migrated probe / llm_api.env stack. -MAPNAV_MODEL = "deepseek-v4-flash" -MAPNAV_EVIDENCE_CHARS = EVIDENCE_TEXT_CHAR_BUDGET -MAPNAV_TOKEN_LIMIT = 100_000 -MAPNAV_PLANNER_THINK_MAX_TOKENS = 16_384 -MAPNAV_TRACE_RAW_CHARS = 2_000 - -# Hard-coded enabled product config (checklist + map). -_PRODUCTION_NAV_DICT: dict[str, Any] = { - "collect_top_k": 6, - "read_score_bonus": 10.0, - "policy": "llm", - "llm_model": MAPNAV_MODEL, - "llm_temperature": 0.0, - "llm_max_tokens": 256, - "planner_llm_max_tokens": 1024, - "harvest_llm_max_tokens": 1024, - "map_char_limit": 5000, - "enable_recursive_dispatch": True, - "max_dispatch_depth": 5, - "subagent_model": MAPNAV_MODEL, - "compose_confidence_weight": 0.5, - "mode": "checklist", - "planning_map_char_limit": 10000, - "planner_max_subgoals": 0, - "planner_model": MAPNAV_MODEL, - "planner_thinking": "enabled", - "planner_think_max_tokens": MAPNAV_PLANNER_THINK_MAX_TOKENS, - "token_limit": MAPNAV_TOKEN_LIMIT, - "subgoal_max_attempts": 2, - "max_replans": 1, - "max_waves": 0, - "max_harvest_depth": 5, - "enable_node_filter": True, - "filter_max_rounds": 3, -} - - -def nav_evidence_chars() -> int: - """Evidence pack budget — code constant.""" - return MAPNAV_EVIDENCE_CHARS - - -def build_nav_config() -> NavConfig: - """Checklist + map production config (enabled items only).""" - cfg = NavConfig.from_dict(dict(_PRODUCTION_NAV_DICT)) - cfg.mode = "checklist" - cfg.policy = "llm" - return cfg diff --git a/deprecated/mapnav/nav_llm_backend.py b/deprecated/mapnav/nav_llm_backend.py deleted file mode 100644 index 60ba68836..000000000 --- a/deprecated/mapnav/nav_llm_backend.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Sync ``nav_chat`` backend for Knowhere map-nav. - -Install once at import (module-level). Forwards ``extra`` / thinking as-is. -Model defaults for empty ``model=`` come from ``nav_config.MAPNAV_MODEL``. -""" - -from __future__ import annotations - -import math -from typing import Any, Dict, List, Optional, Sequence, cast - -from shared.services.retrieval.nav.nav_llm import set_nav_chat_backend -from shared.services.retrieval.nav_config import MAPNAV_MODEL - - -def nav_chat_sync_backend( - *, - messages: Sequence[Dict[str, Any]], - model: str = "", - temperature: float = 0.0, - max_tokens: Optional[int] = None, - response_format: Optional[Dict[str, Any]] = None, - extra: Optional[Dict[str, Any]] = None, - timeout: float = 60.0, - **_kwargs: Any, -) -> Dict[str, Any]: - """Knowhere sync backend matching ``nav_chat``'s injected callable shape.""" - from shared.services.ai.llm_overrides import resolve_text - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - requested = str(model or "").strip() or MAPNAV_MODEL - effective_model, api_key, api_url = resolve_text(requested) - client = get_openai_client( - model=effective_model or requested, - api_key=api_key, - api_url=api_url, - ) - - call_kwargs: Dict[str, Any] = {} - if response_format is not None: - call_kwargs["response_format"] = response_format - - if isinstance(extra, dict): - body = extra.get("extra_body") - if isinstance(body, dict) and body: - call_kwargs["extra_body"] = dict(body) - - timeout_arg: Optional[int] = None - if timeout is not None and math.isfinite(float(timeout)): - timeout_arg = max(1, int(math.ceil(float(timeout)))) - - text, usage = client.chat_completion_with_usage( - cast(List[Any], list(messages)), - model=effective_model or requested, - temperature=float(temperature if temperature is not None else 0.0), - max_tokens=int(max_tokens) if max_tokens is not None else 2048, - timeout=timeout_arg, - **call_kwargs, - ) - return {"content": text or "", "usage": usage or {}} - - -def install_nav_chat_backend() -> None: - """Register the sync backend once (safe to call repeatedly).""" - set_nav_chat_backend(nav_chat_sync_backend) - - -install_nav_chat_backend() diff --git a/deprecated/mapnav/nav_snapshot.py b/deprecated/mapnav/nav_snapshot.py deleted file mode 100644 index 5e1016dd0..000000000 --- a/deprecated/mapnav/nav_snapshot.py +++ /dev/null @@ -1,1307 +0,0 @@ -"""Async namespace snapshot loader for map-nav. - -Loads current-revision sections + chunks for ``(user_id, namespace)``, builds a -synchronous ``NamespaceKnowhereProvider``, and a separate ``chunk_ref_index`` -that keeps **DB-original** ``section_path`` / ``job_id`` (never remounted). -""" - -from __future__ import annotations - -import json -import logging -import os -import time -import zlib -from dataclasses import dataclass -from collections.abc import Callable, Iterator, Mapping -from typing import Any, Protocol - -from sqlalchemy import ( - ARRAY, - Executable, - String, - bindparam, - cast, - func, - literal, - select, - tuple_, -) -from sqlalchemy.engine import Result -from sqlalchemy.exc import SQLAlchemyError - -from shared.models.database.document import ( - Document, - DocumentChunk, - DocumentSection, - RetrievalNamespaceMapSnapshot, - RetrievalNamespaceGeneration, - RetrievalServingRevisionManifest, -) -from shared.models.database.job_result import JobResult -from shared.services.retrieval.nav.nav_knowhere import ( - LazyKnowhereProvider, - NamespaceKnowhereProvider, - ReadOnlyChunkStore, - knowhere_database_url, -) -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, - UnitRow, -) -from shared.services.retrieval.search.section_filters import is_excluded_section -from shared.services.retrieval.serving_manifest import ( - decode_namespace_map_snapshot, - decode_serving_manifest, -) -from shared.services.retrieval.manifest_cache import get_cached_manifest_payloads -from shared.services.retrieval.namespace_map_snapshot_redis import ( - NamespaceMapSnapshotRedisCache, -) -from shared.core.exceptions.redis_exceptions import RedisOperationError - - -# Keep each payload query bounded under the API's 30-second statement timeout. -# Ten-thousand-row keyset pages avoid OFFSET scans while keeping the reference -# payload bounded under the API's 30-second asyncpg command timeout. -_CHUNK_BATCH_SIZE = 10_000 -# Keep revision predicates bounded while reducing round trips for large -# namespaces. Keyset paging still caps each payload query at 10,000 rows. -_REVISION_GROUP_SIZE = 64 -# Compressed manifests are efficient for small namespaces, but transferring a -# large set of binary payloads can exceed the asyncpg statement timeout before -# PostgreSQL has done meaningful work. For larger snapshots, the normalized -# section/chunk index loaders below transfer only the fields map-nav needs and -# page them by revision. -_MANIFEST_MAX_REVISION_COUNT = 32 -_logger = logging.getLogger(__name__) - - -def _is_snapshot_timing_enabled() -> bool: - """Enable detailed snapshot timings for local performance diagnosis.""" - return os.environ.get("RETRIEVAL_SNAPSHOT_TIMING", "0").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -def _count_section_shape( - sections_by_doc: Mapping[str, list[SectionRow]], -) -> tuple[int, int, int, int]: - """Return section rows, distinct paths, roots, and leaves for diagnostics.""" - rows: list[SectionRow] = [ - row for document_rows in sections_by_doc.values() for row in document_rows - ] - parent_ids: set[str] = { - str(row.parent_section_id).strip() - for row in rows - if str(row.parent_section_id or "").strip() - } - return ( - len(rows), - len({str(row.section_path) for row in rows}), - sum(1 for row in rows if not str(row.parent_section_id or "").strip()), - sum(1 for row in rows if str(row.section_id) not in parent_ids), - ) - - -class SnapshotSession(Protocol): - """Minimal database interface required by the snapshot loader.""" - - async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - raise NotImplementedError - - async def rollback(self) -> None: - raise NotImplementedError - - -@dataclass(frozen=True) -class NavSnapshot: - """In-memory corpus for one map-nav episode.""" - - provider: NamespaceKnowhereProvider - chunk_ref_index: Mapping[str, dict[str, Any]] - document_ids: list[str] - document_titles: dict[str, str] - document_revisions: Mapping[str, str] | None - - def close(self) -> None: - close = getattr(self.provider, "close", None) - if callable(close): - close() - - -def build_nav_snapshot( - *, - document_titles: dict[str, str], - sections_by_doc: dict[str, list[SectionRow]], - units_by_doc: dict[str, list[UnitRow]], - chunk_ref_index: dict[str, dict[str, Any]], - document_revisions: Mapping[str, str] | None = None, -) -> NavSnapshot: - """Assemble provider + index from already-fetched rows (also used by tests).""" - doc_ids = [ - did for did in document_titles if did in sections_by_doc or did in units_by_doc - ] - if not doc_ids: - raise ValueError("nav snapshot requires at least one active document") - - providers: list[KnowhereProvider] = [] - for doc_id in doc_ids: - providers.append( - KnowhereProvider( - doc_id=doc_id, - sections=sections_by_doc.get(doc_id, ()), - units=units_by_doc.get(doc_id, ()), - ) - ) - provider = NamespaceKnowhereProvider(providers, titles=document_titles) - return NavSnapshot( - provider=provider, - chunk_ref_index=dict(chunk_ref_index), - document_ids=list(provider.document_ids()), - document_titles={ - did: document_titles.get(did, did) for did in provider.document_ids() - }, - document_revisions=( - { - did: str(document_revisions.get(did, "")) - for did in provider.document_ids() - if document_revisions.get(did) - } - if document_revisions is not None - else None - ), - ) - - -async def load_nav_snapshot( - db: SnapshotSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str] | None = None, - exclude_sections: list[dict[str, str]] | None = None, - lazy: bool = False, - revision_pins: Mapping[str, str] | None = None, - generation: int | None = None, -) -> NavSnapshot: - """Preload namespace current revision into a sync map-nav snapshot.""" - excluded_docs = [ - str(x).strip() for x in (exclude_document_ids or ()) if str(x).strip() - ] - excluded_secs = list(exclude_sections or ()) - snapshot_started = time.perf_counter() - - if revision_pins is None: - doc_stmt = ( - select( - Document.document_id, - Document.source_file_name, - Document.current_job_result_id, - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .where(Document.current_job_result_id.is_not(None)) - .order_by(Document.document_id) - ) - else: - pinned_document_ids = [str(document_id) for document_id in revision_pins] - if not pinned_document_ids: - raise ValueError("revision pins must include at least one document") - doc_stmt = ( - select(Document.document_id, Document.source_file_name) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.document_id.in_(pinned_document_ids)) - .order_by(Document.document_id) - ) - if excluded_docs: - doc_stmt = doc_stmt.where(Document.document_id.notin_(excluded_docs)) - doc_query_started = time.perf_counter() - doc_rows = list((await db.execute(doc_stmt)).all()) - doc_query_seconds = time.perf_counter() - doc_query_started - if not doc_rows: - raise ValueError( - f"no active documents with current revision for " - f"user_id={user_id!r} namespace={namespace!r}" - ) - - document_titles: dict[str, str] = {} - current_job_result_ids: set[str] = set() - document_revisions: list[tuple[str, str]] = [] - for doc_row in doc_rows: - document_id = doc_row[0] - source_file_name = doc_row[1] - did = str(document_id) - title = str(source_file_name or "").strip() or did - document_titles[did] = title - job_result_id = ( - str(revision_pins.get(did, "")) - if revision_pins is not None - else str(doc_row[2] or "") - ) - if job_result_id: - current_job_result_ids.add(job_result_id) - document_revisions.append((did, job_result_id)) - - job_query_started = time.perf_counter() - job_result_rows = await db.execute( - select(JobResult.id, JobResult.job_id).where( - JobResult.id.in_(list(current_job_result_ids)) - ) - ) - job_query_seconds = time.perf_counter() - job_query_started - 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 - } - - snapshot_entries = await _resolve_namespace_snapshot_entries( - db, - user_id=user_id, - namespace=namespace, - document_revisions=document_revisions, - expected_generation=generation, - ) - if snapshot_entries is not None: - parse_started = time.perf_counter() - manifest_sections = _parse_manifest_entries( - snapshot_entries, - exclude_sections=excluded_secs, - job_id_by_result_id=job_id_by_result_id, - ) - parse_seconds = time.perf_counter() - parse_started - else: - _logger.warning( - "retrieval snapshot fallback=manifest_merge user_id=%s namespace=%s documents=%d", - user_id, - namespace, - len(document_revisions), - ) - manifest_sections = None - if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT: - manifest_sections = await _load_manifest_sections( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - job_id_by_result_id=job_id_by_result_id, - ) - parse_seconds = 0.0 - if manifest_sections is None: - _logger.warning( - "retrieval snapshot fallback=table_scan user_id=%s namespace=%s documents=%d", - user_id, - namespace, - len(document_revisions), - ) - sections_by_doc, section_path_by_id = await _load_sections( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - ) - manifest_chunk_index = None - else: - sections_by_doc, section_path_by_id, manifest_chunk_index = manifest_sections - if lazy: - if manifest_chunk_index is None: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = await _load_chunk_index( - db, - 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, - ) - else: - chunk_ids_by_doc, chunk_ref_index, remounted_assets = manifest_chunk_index - kept_titles = { - did: title - for did, title in document_titles.items() - if sections_by_doc.get(did) - } - if not kept_titles: - raise ValueError( - f"nav snapshot empty after excludes for " - f"user_id={user_id!r} namespace={namespace!r}" - ) - revisions = { - did: result_id - for did, result_id in document_revisions - if did in kept_titles - } - store = ReadOnlyChunkStore( - dsn=knowhere_database_url(), - revisions=revisions, - excluded_sections={ - ( - str(item.get("document_id") or "").strip(), - section_id, - ) - for section_id, section_path in section_path_by_id.items() - for item in excluded_secs - if isinstance(item, dict) - and str(item.get("document_id") or "").strip() - and str(item.get("section_path") or "").strip() == section_path - }, - ) - try: - providers = [ - LazyKnowhereProvider( - doc_id=did, - sections=sections_by_doc.get(did, ()), - chunk_store=store, - known_chunk_ids=chunk_ids_by_doc.get(did, ()), - root_asset_ids=remounted_assets.get(did, {}).get("root", ()), - remounted_assets_by_section=remounted_assets.get(did, {}).get( - "owners", {} - ), - ) - for did in kept_titles - ] - provider = NamespaceKnowhereProvider( - providers, - titles=kept_titles, - chunk_owner_by_id={ - chunk_id: document_id - for document_id, chunk_ids in chunk_ids_by_doc.items() - for chunk_id in chunk_ids - }, - ) - ref_index_started = time.perf_counter() - lazy_ref_index = LazyChunkRefIndex( - chunk_ref_index, - resolver=store.load_chunk_reference_metadata, - ) - ref_index_seconds = time.perf_counter() - ref_index_started - except Exception: - store.close() - raise - snapshot = NavSnapshot( - provider=provider, - chunk_ref_index=lazy_ref_index, - document_ids=list(provider.document_ids()), - document_titles={ - did: kept_titles.get(did, did) for did in provider.document_ids() - }, - document_revisions=dict(revisions), - ) - if _is_snapshot_timing_enabled(): - section_rows, section_paths, root_sections, leaf_sections = ( - _count_section_shape(sections_by_doc) - ) - _logger.info( - "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " - "ref_index_copy_seconds=%.3f doc_query_seconds=%.3f " - "job_query_seconds=%.3f documents=%d section_rows=%d " - "section_paths=%d root_sections=%d leaf_sections=%d refs=%d " - "mode=lazy", - time.perf_counter() - snapshot_started, - parse_seconds, - ref_index_seconds, - doc_query_seconds, - job_query_seconds, - len(snapshot.document_ids), - section_rows, - section_paths, - root_sections, - leaf_sections, - len(snapshot.chunk_ref_index), - ) - return snapshot - - units_by_doc, chunk_ref_index = await _load_chunks( - db, - 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. - kept_titles = { - did: title for did, title in document_titles.items() if sections_by_doc.get(did) - } - if not kept_titles: - raise ValueError( - f"nav snapshot empty after excludes for " - f"user_id={user_id!r} namespace={namespace!r}" - ) - - snapshot = build_nav_snapshot( - document_titles=kept_titles, - sections_by_doc={did: sections_by_doc.get(did, []) for did in kept_titles}, - units_by_doc={did: units_by_doc.get(did, []) for did in kept_titles}, - chunk_ref_index=chunk_ref_index, - document_revisions={ - document_id: job_result_id - for document_id, job_result_id in document_revisions - if document_id in kept_titles - }, - ) - if _is_snapshot_timing_enabled(): - section_rows, section_paths, root_sections, leaf_sections = ( - _count_section_shape(sections_by_doc) - ) - _logger.info( - "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " - "doc_query_seconds=%.3f job_query_seconds=%.3f documents=%d " - "section_rows=%d section_paths=%d root_sections=%d " - "leaf_sections=%d refs=%d mode=eager", - time.perf_counter() - snapshot_started, - parse_seconds, - doc_query_seconds, - job_query_seconds, - len(snapshot.document_ids), - section_rows, - section_paths, - root_sections, - leaf_sections, - len(snapshot.chunk_ref_index), - ) - return snapshot - - -async def _resolve_namespace_snapshot_entries( - db: SnapshotSession, - *, - user_id: str, - namespace: str, - document_revisions: list[tuple[str, str]], - expected_generation: int | None = None, -) -> list[tuple[object, ...]] | None: - """Return manifest-shaped entries from the persisted namespace snapshot. - - Returns ``None`` (triggering the exact per-revision fallback) when the - snapshot row is missing, corrupt, or stale for any requested revision. - """ - generation_statement = select(RetrievalNamespaceGeneration.generation).where( - RetrievalNamespaceGeneration.user_id == user_id, - RetrievalNamespaceGeneration.namespace == namespace, - ) - try: - generation_result = await db.execute(generation_statement) - _current_generation = generation_result.scalar_one_or_none() - except SQLAlchemyError as exc: - await db.rollback() - _logger.warning( - "retrieval snapshot generation lookup failed; using fallback error=%s", - exc, - ) - return None - if _is_snapshot_timing_enabled(): - _logger.info( - "retrieval snapshot generation generation=%s", - _current_generation, - ) - generation_value = int(expected_generation) if expected_generation is not None else None - statement = select( - RetrievalNamespaceMapSnapshot.generation, - RetrievalNamespaceMapSnapshot.checksum, - RetrievalNamespaceMapSnapshot.format_version, - ).where( - RetrievalNamespaceMapSnapshot.user_id == user_id, - RetrievalNamespaceMapSnapshot.namespace == namespace, - ) - lookup_started = time.perf_counter() - try: - row = (await db.execute(statement)).first() - except SQLAlchemyError as exc: - await db.rollback() - _logger.warning( - "retrieval snapshot namespace lookup failed user_id=%s namespace=%s error=%s", - user_id, - namespace, - exc, - ) - return None - lookup_seconds = time.perf_counter() - lookup_started - if _is_snapshot_timing_enabled(): - _logger.info( - "retrieval snapshot lookup seconds=%.3f found=%s", - lookup_seconds, - row is not None, - ) - if row is None: - return None - generation, checksum, format_version = row - row_generation = int(generation) - if generation_value is not None and row_generation != generation_value: - _logger.info( - "retrieval snapshot generation mismatch row=%d expected=%d", - row_generation, - generation_value, - ) - return None - if ( - generation_value is None - and _current_generation is not None - and int(_current_generation) > 0 - and row_generation != int(_current_generation) - ): - _logger.info( - "retrieval snapshot generation mismatch row=%d current=%d", - row_generation, - int(_current_generation), - ) - return None - cached_blob: bytes | None = None - try: - cached_blob = await NamespaceMapSnapshotRedisCache.get( - user_id=user_id, namespace=namespace, generation=row_generation - ) - except RedisOperationError as exc: - _logger.warning("retrieval snapshot redis get failed error=%s", exc) - database_blob: bytes | None = None - if cached_blob is None: - payload_result = await db.execute( - select(RetrievalNamespaceMapSnapshot.payload_zlib).where( - RetrievalNamespaceMapSnapshot.user_id == user_id, - RetrievalNamespaceMapSnapshot.namespace == namespace, - ) - ) - payload_row = payload_result.first() - if payload_row is None or payload_row[0] is None: - return None - database_blob = bytes(payload_row[0]) - blob: bytes = database_blob - else: - blob = cached_blob - is_timing_enabled = _is_snapshot_timing_enabled() - decode_timings: dict[str, float] | None = {} if is_timing_enabled else None - decode_started = time.perf_counter() - try: - payload = decode_namespace_map_snapshot( - blob, - checksum=str(checksum), - format_version=int(format_version), - timings=decode_timings, - ) - except (ValueError, TypeError, zlib.error): - if cached_blob is not None: - # A stale/corrupt cache entry must never shadow the PostgreSQL source. - try: - if database_blob is None: - fallback_result = await db.execute( - select(RetrievalNamespaceMapSnapshot.payload_zlib).where( - RetrievalNamespaceMapSnapshot.user_id == user_id, - RetrievalNamespaceMapSnapshot.namespace == namespace, - ) - ) - fallback_row = fallback_result.first() - if fallback_row is None or fallback_row[0] is None: - return None - database_blob = bytes(fallback_row[0]) - assert database_blob is not None - payload = decode_namespace_map_snapshot( - database_blob, - checksum=str(checksum), - format_version=int(format_version), - timings=decode_timings, - ) - blob = database_blob - cached_blob = None - except (ValueError, TypeError, zlib.error): - return None - else: - return None - decoded_documents = payload.get("documents") - if not isinstance(decoded_documents, dict): - return None - documents = decoded_documents - if cached_blob is None: - try: - await NamespaceMapSnapshotRedisCache.set( - user_id=user_id, - namespace=namespace, - generation=row_generation, - payload_zlib=blob, - ) - except RedisOperationError as exc: - _logger.warning("retrieval snapshot redis set failed error=%s", exc) - if is_timing_enabled: - _logger.info( - "retrieval snapshot decode cache_hit=false seconds=%.3f " - "compressed_bytes=%d decompressed_bytes=%d " - "decompress_seconds=%.3f checksum_seconds=%.3f " - "json_decode_seconds=%.3f documents=%d", - time.perf_counter() - decode_started, - int((decode_timings or {}).get("compressed_bytes", 0.0)), - int((decode_timings or {}).get("decompressed_bytes", 0.0)), - (decode_timings or {}).get("decompress_seconds", 0.0), - (decode_timings or {}).get("checksum_seconds", 0.0), - (decode_timings or {}).get("json_decode_seconds", 0.0), - len(documents), - ) - elif is_timing_enabled: - _logger.info( - "retrieval snapshot decode cache_hit=true seconds=%.3f " - "compressed_bytes=%d decompressed_bytes=%d decompress_seconds=%.3f " - "checksum_seconds=%.3f json_decode_seconds=%.3f documents=%d", - time.perf_counter() - decode_started, - int((decode_timings or {}).get("compressed_bytes", 0.0)), - int((decode_timings or {}).get("decompressed_bytes", 0.0)), - (decode_timings or {}).get("decompress_seconds", 0.0), - (decode_timings or {}).get("checksum_seconds", 0.0), - (decode_timings or {}).get("json_decode_seconds", 0.0), - len(documents), - ) - entries: list[tuple[object, ...]] = [] - for document_id, job_result_id in document_revisions: - entry = documents.get(document_id) - if not isinstance(entry, dict) or str(entry.get("job_result_id") or "") != job_result_id: - return None - entries.append((document_id, job_result_id, entry, None, None)) - return entries - - -async def _resolve_manifest_entries( - db: SnapshotSession, - *, - document_revisions: list[tuple[str, str]], -) -> list[tuple[object, ...]] | None: - """Resolve per-revision manifest rows from the request cache or table.""" - cached_payloads = get_cached_manifest_payloads( - db, - revisions=dict(document_revisions), - ) - if cached_payloads is not None and len(cached_payloads) == len(document_revisions): - manifest_entries: list[tuple[object, ...]] = [ - (document_id, job_result_id, payload, None, None) - for document_id, job_result_id in document_revisions - for payload in [cached_payloads.get((document_id, job_result_id))] - if payload is not None - ] - if len(manifest_entries) != len(document_revisions): - return None - return manifest_entries - - manifest_entries = [] - for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): - revision_group = document_revisions[ - group_start : group_start + _REVISION_GROUP_SIZE - ] - statement = select( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - ).in_(revision_group) - ) - try: - rows = (await db.execute(statement)).all() - except SQLAlchemyError as exc: - await db.rollback() - _logger.warning( - "retrieval manifest lookup failed revisions=%s error=%s", - revision_group, - exc, - ) - return None - if len(rows) != len(revision_group): - return None - manifest_entries.extend(tuple(row) for row in rows) - return manifest_entries - - -async def _load_manifest_sections( - db: SnapshotSession, - *, - document_revisions: list[tuple[str, str]], - exclude_sections: list[dict[str, str]], - job_id_by_result_id: dict[str, str], -) -> tuple[ - dict[str, list[SectionRow]], - dict[str, str], - tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], -] | None: - """Load section metadata from complete serving manifests when available.""" - manifest_entries = await _resolve_manifest_entries( - db, document_revisions=document_revisions - ) - if manifest_entries is None: - return None - return _parse_manifest_entries( - manifest_entries, - exclude_sections=exclude_sections, - job_id_by_result_id=job_id_by_result_id, - ) - - -def _parse_manifest_entries( - manifest_entries: list[tuple[object, ...]], - *, - exclude_sections: list[dict[str, str]], - job_id_by_result_id: dict[str, str], -) -> tuple[ - dict[str, list[SectionRow]], - dict[str, str], - tuple[dict[str, list[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]], -] | None: - """Parse manifest-shaped entries (decoded dict or raw compressed row) into sections.""" - by_doc: dict[str, list[SectionRow]] = {} - path_by_id: dict[str, str] = {} - ids_by_doc: dict[str, list[str]] = {} - ref_index: dict[str, dict[str, Any]] = {} - root_assets_by_doc: dict[str, set[str]] = {} - text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} - section_seconds = 0.0 - chunk_seconds = 0.0 - section_count = 0 - chunk_count = 0 - connection_count = 0 - try: - for document_id, _job_result_id, payload_zlib, checksum, format_version in manifest_entries: - if isinstance(payload_zlib, dict): - payload = payload_zlib - else: - if not isinstance(payload_zlib, (bytes, bytearray, memoryview)): - return None - if checksum is None or format_version is None: - return None - payload = decode_serving_manifest( - bytes(payload_zlib), - checksum=str(checksum), - format_version=int(str(format_version)), - ) - section_started = time.perf_counter() - raw_sections = payload.get("sections") - if not isinstance(raw_sections, list): - return None - for raw_section in raw_sections: - if not isinstance(raw_section, dict): - return None - section_path = str(raw_section.get("section_path") or "") - if is_excluded_section( - document_id=str(document_id), - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - section_id = str(raw_section.get("section_id") or "") - if not section_id or not section_path: - return None - section = SectionRow( - section_id=section_id, - parent_section_id=( - str(raw_section["parent_section_id"]) - if raw_section.get("parent_section_id") - else None - ), - section_path=section_path, - section_title=str(raw_section.get("section_title") or "").strip(), - section_level=int(raw_section.get("section_level") or 0), - summary=str(raw_section.get("summary") or "").strip(), - sort_order=int(raw_section.get("sort_order") or 0), - ) - by_doc.setdefault(str(document_id), []).append(section) - path_by_id[section_id] = section_path - section_count += 1 - section_seconds += time.perf_counter() - section_started - chunk_started = time.perf_counter() - raw_chunks = payload.get("chunks") - if not isinstance(raw_chunks, list): - return None - for raw_chunk in raw_chunks: - if not isinstance(raw_chunk, dict): - return None - chunk_id = str(raw_chunk.get("chunk_id") or "").strip() - if not chunk_id: - return None - section_id = ( - str(raw_chunk["section_id"]) - if raw_chunk.get("section_id") - else None - ) - section_path = path_by_id.get(section_id) if section_id else None - if is_excluded_section( - document_id=str(document_id), - section_path=section_path, - exclude_sections=exclude_sections, - ) or (section_id and section_id not in path_by_id): - continue - chunk_type = str(raw_chunk.get("chunk_type") or "text") - meta = { - "document_id": str(document_id), - "section_path": section_path, - "chunk_type": chunk_type, - "file_path": None, - "job_id": job_id_by_result_id.get(str(_job_result_id)), - } - document_key = str(document_id) - ids_by_doc.setdefault(document_key, []).append(chunk_id) - ref_index[chunk_id] = meta - ref_index[f"{document_key}:{chunk_id}"] = meta - if ( - chunk_type in {"image", "table"} - and section_id - and section_path == "Root" - ): - root_assets_by_doc.setdefault(document_key, set()).add(chunk_id) - connections = raw_chunk.get("connect_to") - if chunk_type == "text" and isinstance(connections, list): - for connection in connections: - target = ( - str(connection.get("target") or "").strip() - if isinstance(connection, dict) - else str(connection or "").strip() - ) - if target: - text_connections_by_doc.setdefault(document_key, []).append( - (section_id or "", target) - ) - connection_count += 1 - chunk_count += 1 - chunk_seconds += time.perf_counter() - chunk_started - except (TypeError, ValueError, KeyError): - return None - remounted: dict[str, dict[str, Any]] = {} - remount_started = time.perf_counter() - for document_id, asset_ids in root_assets_by_doc.items(): - owners: dict[str, list[str]] = {} - for section_id, target in text_connections_by_doc.get(document_id, ()): - if target in asset_ids: - owners.setdefault(section_id, []).append(target) - remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} - if _is_snapshot_timing_enabled(): - _logger.info( - "retrieval snapshot parse sections=%d chunks=%d connections=%d " - "section_seconds=%.3f chunk_seconds=%.3f remount_seconds=%.3f " - "ref_index_keys=%d", - section_count, - chunk_count, - connection_count, - section_seconds, - chunk_seconds, - time.perf_counter() - remount_started, - len(ref_index), - ) - return by_doc, path_by_id, (ids_by_doc, ref_index, remounted) - - -async def _load_chunk_index( - db: SnapshotSession, - *, - 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[str]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]: - """Load only IDs/reference metadata; content and asset paths remain lazy.""" - ids_by_doc: dict[str, list[str]] = {} - ref_index: dict[str, dict[str, Any]] = {} - root_assets_by_doc: dict[str, set[str]] = {} - text_connections_by_doc: dict[str, list[tuple[str, str]]] = {} - query_seconds = 0.0 - assembly_seconds = 0.0 - query_count = 0 - 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, - # Asset paths are only needed for selected references and - # are resolved by ``LazyChunkRefIndex`` at bridge time. - DocumentChunk.chunk_metadata["connect_to"].label("connect_to"), - DocumentChunk.sort_order, - DocumentChunk.id, - ) - .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, - ) - .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(value) for value in last_key]) - ) - query_started = time.perf_counter() - rows = (await db.execute(stmt)).all() - query_seconds += time.perf_counter() - query_started - query_count += 1 - if not rows: - break - assembly_started = time.perf_counter() - 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 = 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, - ) or (section_id and section_id not in section_path_by_id): - continue - if not chunk_id: - continue - chunk_type = str(row[4] or "text") - meta = { - "document_id": document_id, - "section_path": section_path, - "chunk_type": chunk_type, - "file_path": None, - "job_id": job_id_by_result_id.get(job_result_id), - } - ids_by_doc.setdefault(document_id, []).append(chunk_id) - ref_index[chunk_id] = meta - ref_index[f"{document_id}:{chunk_id}"] = meta - if ( - chunk_type in {"image", "table"} - and section_id - and section_path == "Root" - ): - root_assets_by_doc.setdefault(document_id, set()).add(chunk_id) - connections = row[5] - if isinstance(connections, str) and connections.strip(): - try: - connections = json.loads(connections) - except json.JSONDecodeError: - connections = None - if chunk_type == "text" and isinstance(connections, list): - for connection in connections: - if not isinstance(connection, dict): - continue - target = str(connection.get("target") or "").strip() - if not target: - continue - text_connections_by_doc.setdefault(document_id, []).append( - (section_id or "", target) - ) - assembly_seconds += time.perf_counter() - assembly_started - last = rows[-1] - last_key = ( - str(last[0]), - str(last[1]), - int(last[6] or 0), - str(last[2] or ""), - str(last[7]), - ) - if len(rows) < _CHUNK_BATCH_SIZE: - break - remounted: dict[str, dict[str, Any]] = {} - for document_id, asset_ids in root_assets_by_doc.items(): - owners: dict[str, list[str]] = {} - for section_id, target in text_connections_by_doc.get(document_id, ()): - if target in asset_ids: - owners.setdefault(section_id, []).append(target) - remounted[document_id] = {"root": sorted(asset_ids), "owners": owners} - _logger.info( - "retrieval snapshot phase=chunk_index rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", - sum(len(chunk_ids) for chunk_ids in ids_by_doc.values()), - query_count, - query_seconds, - assembly_seconds, - ) - return ids_by_doc, ref_index, remounted - - -class LazyChunkRefIndex(Mapping[str, dict[str, Any]]): - """Reference metadata map that resolves selected asset paths on demand. - - Snapshot construction still records every chunk identity, section path, - type, and job id so navigation ownership is unchanged. ``file_path`` is - fetched only when the exit bridge asks for a selected reference. - """ - - def __init__( - self, - base: Mapping[str, dict[str, Any]], - *, - resolver: Callable[[str, str], Mapping[str, Any] | None], - ) -> None: - self._base = {str(key): dict(value) for key, value in base.items()} - self._resolver = resolver - - def __getitem__(self, key: str) -> dict[str, Any]: - value = self.get(key) - if value is None: - raise KeyError(key) - return value - - def __iter__(self) -> Iterator[str]: - return iter(self._base) - - def __len__(self) -> int: - return len(self._base) - - def get(self, key: str, default: Any = None) -> dict[str, Any] | Any: - value = self._base.get(str(key)) - if value is None: - return default - if value.get("file_path") is None: - document_id = str(value.get("document_id") or "").strip() - chunk_id = str(key).split(":", 1)[-1].strip() - resolved = self._resolver(document_id, chunk_id) - if resolved is not None: - value.update({"file_path": resolved.get("file_path") or None}) - return value - - -async def _load_sections( - db: SnapshotSession, - *, - 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. - by_doc: dict[str, list[SectionRow]] = {} - path_by_id: dict[str, str] = {} - document_ids = [document_id for document_id, _ in document_revisions] - job_result_ids = [job_result_id for _, job_result_id in document_revisions] - revision_rows = ( - func.unnest( - cast( - bindparam("section_document_ids", value=document_ids), ARRAY(String()) - ), - cast( - bindparam("section_job_result_ids", value=job_result_ids), - ARRAY(String()), - ), - ) - .table_valued("document_id", "job_result_id") - .render_derived(name="revisions") - ) - last_key: tuple[str, str, int, str] | None = None - query_seconds = 0.0 - assembly_seconds = 0.0 - query_count = 0 - row_count = 0 - while True: - stmt = ( - select( - DocumentSection.document_id, - DocumentSection.section_id, - DocumentSection.parent_section_id, - DocumentSection.section_path, - DocumentSection.section_title, - DocumentSection.section_level, - DocumentSection.summary, - DocumentSection.sort_order, - DocumentSection.job_result_id, - ) - .join( - revision_rows, - (DocumentSection.document_id == revision_rows.c.document_id) - & (DocumentSection.job_result_id == revision_rows.c.job_result_id), - ) - .order_by( - DocumentSection.document_id, - DocumentSection.job_result_id, - DocumentSection.sort_order, - DocumentSection.section_id, - ) - .limit(_CHUNK_BATCH_SIZE) - ) - if last_key is not None: - stmt = stmt.where( - tuple_( - DocumentSection.document_id, - DocumentSection.job_result_id, - DocumentSection.sort_order, - DocumentSection.section_id, - ) - > tuple_(*[literal(value) for value in last_key]) - ) - query_started = time.perf_counter() - rows = (await db.execute(stmt)).all() - query_seconds += time.perf_counter() - query_started - query_count += 1 - if not rows: - break - row_count += len(rows) - assembly_started = time.perf_counter() - for row in rows: - document_id = str(row[0]) - section_path = str(row[3] or "") - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - section_id = str(row[1]) - section = SectionRow( - section_id=section_id, - parent_section_id=str(row[2]) if row[2] else None, - section_path=section_path, - section_title=str(row[4] or "").strip(), - section_level=int(row[5] or 0), - summary=str(row[6] or "").strip(), - sort_order=int(row[7] or 0), - ) - by_doc.setdefault(document_id, []).append(section) - path_by_id[section_id] = section_path - assembly_seconds += time.perf_counter() - assembly_started - last = rows[-1] - last_key = ( - str(last[0]), - str(last[8]), - int(last[7] or 0), - str(last[1]), - ) - if len(rows) < _CHUNK_BATCH_SIZE: - break - _logger.info( - "retrieval snapshot phase=sections rows=%d queries=%d query_seconds=%.3f assembly_seconds=%.3f", - row_count, - query_count, - query_seconds, - assembly_seconds, - ) - return by_doc, path_by_id - - -async def _load_chunks( - db: SnapshotSession, - *, - 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]]]: - # 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 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, - ) - .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, - ) - .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 - - 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]), - ) - if len(rows) < _CHUNK_BATCH_SIZE: - break - return by_doc, ref_index - - -def _as_meta(value: Any) -> dict[str, Any]: - if isinstance(value, dict): - return dict(value) - return {} diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_snapshot_quality_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_snapshot_quality_contract.py deleted file mode 100644 index 00855e6b5..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_snapshot_quality_contract.py +++ /dev/null @@ -1,294 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from collections.abc import Sequence -from collections import Counter -import math -from typing import Any - -from shared.services.retrieval.nav.nav_knowhere import ( - LazyKnowhereProvider, - NamespaceKnowhereProvider, - knowhere_database_url, -) -from shared.services.retrieval.nav.nav_map_scores import ( - compute_corpus_map_and_unit_scores, - compute_corpus_map_and_unit_scores_many, -) -from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_hybrid import ( - PersistedBm25Stats, - PersistedScoreCorpus, - PersistedScoreUnit, - score_persisted_corpus_many, -) -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, - UnitRow, -) -from shared.services.retrieval.scoring.score_units import build_score_units - - -@dataclass -class _FakeChunkStore: - units_by_section: dict[str, list[UnitRow]] - section_loads: int = 0 - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> list[UnitRow]: - del document_id - self.section_loads += 1 - units = list(self.units_by_section.get(section_id, ())) - known = {unit.chunk_id for unit in units} - for units_in_section in self.units_by_section.values(): - for unit in units_in_section: - if unit.chunk_id in extra_chunk_ids and unit.chunk_id not in known: - units.append(unit) - return units - - def close(self) -> None: - return None - - -def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace, _FakeChunkStore]: - sections = [ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("section", "root", "Root / Section", "Section", 1, "", 1), - SectionRow("leaf", "section", "Root / Section / Leaf", "Leaf", 2, "", 2), - ] - text = UnitRow( - "duplicate-chunk", - "leaf", - "text", - "alpha retrieval evidence", - 1, - metadata={"connect_to": [{"target": "asset-1", "relation": "embeds"}]}, - ) - asset = UnitRow( - "asset-1", - "root", - "image", - "", - 2, - file_path="images/asset.png", - metadata={"summary": "supporting image"}, - ) - eager = NamespaceKnowhereProvider( - [KnowhereProvider(doc_id="doc", sections=sections, units=[text, asset])], - titles={"doc": "document"}, - ) - store = _FakeChunkStore({"root": [asset], "leaf": [text]}) - lazy = NamespaceKnowhereProvider( - [ - LazyKnowhereProvider( - doc_id="doc", - sections=sections, - chunk_store=store, - known_chunk_ids=[text.chunk_id, asset.chunk_id], - root_asset_ids=[asset.chunk_id], - remounted_assets_by_section={"leaf": [asset.chunk_id]}, - ) - ], - titles={"doc": "document"}, - chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"}, - ) - return ProviderToolSpace(eager), ProviderToolSpace(lazy), store - - -def _multi_document_providers() -> tuple[ - ProviderToolSpace, - ProviderToolSpace, - _FakeChunkStore, -]: - first_sections = [ - SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), - SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), - ] - second_sections = [ - SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), - SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), - ] - first_unit = UnitRow("chunk-a", "leaf-a", "text", "alpha evidence", 1) - second_unit = UnitRow("chunk-b", "leaf-b", "text", "beta evidence", 1) - eager = NamespaceKnowhereProvider( - [ - KnowhereProvider( - doc_id="doc-a", - sections=first_sections, - units=[first_unit], - ), - KnowhereProvider( - doc_id="doc-b", - sections=second_sections, - units=[second_unit], - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - ) - store = _FakeChunkStore( - { - "leaf-a": [first_unit], - "leaf-b": [second_unit], - } - ) - lazy = NamespaceKnowhereProvider( - [ - LazyKnowhereProvider( - doc_id="doc-a", - sections=first_sections, - chunk_store=store, - known_chunk_ids=[first_unit.chunk_id], - ), - LazyKnowhereProvider( - doc_id="doc-b", - sections=second_sections, - chunk_store=store, - known_chunk_ids=[second_unit.chunk_id], - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - chunk_owner_by_id={"chunk-a": "doc-a", "chunk-b": "doc-b"}, - ) - return ProviderToolSpace(eager), ProviderToolSpace(lazy), store - - -def test_lazy_provider_preserves_score_units_and_scores() -> None: - eager, lazy, store = _providers() - - assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") - eager_scores = compute_corpus_map_and_unit_scores( - eager, doc_ids=["doc"], query="alpha retrieval" - ) - lazy_scores = compute_corpus_map_and_unit_scores( - lazy, doc_ids=["doc"], query="alpha retrieval" - ) - assert eager_scores == lazy_scores - - lazy_provider = lazy._provider - self_units = getattr(lazy_provider, "self_units") - assert [unit.chunk_id for unit in self_units("leaf")] == [ - "duplicate-chunk", - "asset-1", - ] - - -def test_persisted_score_projection_ranks_matching_units() -> None: - rows: list[dict[str, str]] = [ - { - "chunk_id": "unit-a", - "path_search_text": "root alpha", - "content_search_text": "common alpha alpha evidence", - }, - { - "chunk_id": "unit-b", - "path_search_text": "root beta", - "content_search_text": "common beta evidence", - }, - { - "chunk_id": "unit-c", - "path_search_text": "root common", - "content_search_text": "common evidence", - }, - ] - queries = ["common alpha", "beta evidence"] - query_tokens = {token for query in queries for token in query.split()} - - def build_stats(search_field: str) -> PersistedBm25Stats: - token_rows = [str(row[search_field]).split() for row in rows] - document_frequency = Counter( - token for tokens in token_rows for token in set(tokens) - ) - document_count = len(token_rows) - raw_idfs = [ - math.log(document_count - frequency + 0.5) - math.log(frequency + 0.5) - for frequency in document_frequency.values() - ] - return PersistedBm25Stats( - document_count=document_count, - total_length=sum(len(tokens) for tokens in token_rows), - document_frequency={ - token: document_frequency[token] for token in query_tokens - }, - average_idf=sum(raw_idfs) / len(raw_idfs), - ) - - corpus = PersistedScoreCorpus( - units=[ - PersistedScoreUnit( - unit_id=str(row["chunk_id"]), - path_length=len(str(row["path_search_text"]).split()), - content_length=len(str(row["content_search_text"]).split()), - path_frequencies={ - token: str(row["path_search_text"]).split().count(token) - for token in query_tokens - }, - content_frequencies={ - token: str(row["content_search_text"]).split().count(token) - for token in query_tokens - }, - ) - for row in rows - ], - path_stats=build_stats("path_search_text"), - content_stats=build_stats("content_search_text"), - ) - - scored = score_persisted_corpus_many(corpus, queries) - assert set(scored) == set(queries) - assert scored["common alpha"]["unit-a"] > scored["common alpha"]["unit-b"] - assert scored["beta evidence"]["unit-b"] > scored["beta evidence"]["unit-a"] - - -def test_missing_index_falls_back_to_legacy_payload_scoring() -> None: - _eager, lazy, store = _providers() - queries: list[str] = ["alpha retrieval", "supporting image"] - store.section_loads = 0 - actual = compute_corpus_map_and_unit_scores_many( - lazy, - doc_ids=["doc"], - queries=queries, - ) - - assert set(actual) == set(queries) - assert any(unit_scores for _map_scores, unit_scores in actual.values()) - assert store.section_loads > 0 - - -def test_missing_index_is_empty_across_documents() -> None: - eager, lazy, store = _multi_document_providers() - queries: list[str] = ["alpha evidence", "beta evidence"] - store.section_loads = 0 - expected = compute_corpus_map_and_unit_scores_many( - eager, - doc_ids=["doc-a", "doc-b"], - queries=queries, - ) - actual = compute_corpus_map_and_unit_scores_many( - lazy, - doc_ids=["doc-a", "doc-b"], - queries=queries, - ) - - assert actual == expected - assert any(unit_scores for _map_scores, unit_scores in actual.values()) - assert store.section_loads > 0 - - -def test_native_chunk_store_strips_async_driver_from_database_url( - monkeypatch: Any, -) -> None: - monkeypatch.setenv( - "DATABASE_URL", - "postgresql+asyncpg://prod-user:prod-password@db.example/knowhere", - ) - monkeypatch.delenv("KNOWHERE_DATABASE_URL", raising=False) - - assert ( - knowhere_database_url() - == "postgresql://prod-user:prod-password@db.example/knowhere" - ) diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_tree_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_tree_contract.py deleted file mode 100644 index 038756a95..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_tree_contract.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Contract coverage for lazy MAP-NAV tree traversal.""" - -from __future__ import annotations - -from shared.services.retrieval.scoring.hierarchy import NodeMeta, ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_provider import KnowhereProvider, SectionRow -from shared.services.retrieval.scoring.score_units import _walk_tree -from shared.services.retrieval.nav._compat import Chunk -from shared.services.retrieval.nav.nav_compose import pack_nav_evidence -from shared.services.retrieval.nav.nav_types import NavConfig, NavState - - -class _MetadataForbiddenProvider(KnowhereProvider): - def node_meta(self, section_id: str) -> NodeMeta: - raise AssertionError(f"tree traversal materialized metadata for {section_id}") - - -def test_tree_walk_reads_children_and_titles_without_materializing_metadata() -> None: - provider = _MetadataForbiddenProvider( - doc_id="doc", - sections=[ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), - ], - units=(), - ) - - children, leaves, titles = _walk_tree( - ProviderToolSpace(provider), - "doc", - ["root"], - ) - - assert children == {"root": ["child"], "child": []} - assert leaves == {"child"} - assert titles == {"root": "Root", "child": "Child"} - - -def test_evidence_pack_reads_titles_without_materializing_subtree_metadata() -> None: - class _CountingMetadataProvider(_MetadataForbiddenProvider): - metadata_calls = 0 - - def node_meta(self, section_id: str) -> NodeMeta: - self.metadata_calls += 1 - return super().node_meta(section_id) - - provider = _CountingMetadataProvider( - doc_id="doc", - sections=[ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), - ], - units=(), - ) - toolspace = ProviderToolSpace(provider) - state = NavState(doc_id="doc", query="child") - chunk = Chunk( - node_id="child", - doc_id="doc", - text="evidence", - line_ids=(1,), - section_id="child", - ) - - result = pack_nav_evidence( - [(chunk, 1.0)], - toolspace, - state, - NavConfig(), - budget_chars=100, - ) - - assert result.evidence_text == "[E1]\n[§ Root / Child]\nevidence" - assert provider.metadata_calls == 0 - - -def test_evidence_pack_identifies_header_owners_from_parent_chain() -> None: - provider = KnowhereProvider( - doc_id="doc", - sections=[ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("parent", "root", "Root / Parent", "Parent", 1, "", 1), - SectionRow( - "child", "parent", "Root / Parent / Child", "Child", 2, "", 2 - ), - ], - units=(), - ) - toolspace = ProviderToolSpace(provider) - state = NavState(doc_id="doc", query="child") - chunks = [ - Chunk("parent", "doc", "parent evidence", (1,), "parent"), - Chunk("child", "doc", "child evidence", (2,), "child"), - ] - - result = pack_nav_evidence( - [(chunks[0], 1.0), (chunks[1], 0.9)], - toolspace, - state, - NavConfig(), - budget_chars=200, - ) - - assert result.kept_chunks == [chunks[1]] - assert result.evidence_text == "[E1]\n[§ Root / Parent / Child]\nchild evidence" diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_map_score_parity_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_map_score_parity_contract.py deleted file mode 100644 index c7d651c07..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_map_score_parity_contract.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Contract tests for map score pooling semantics.""" - -from __future__ import annotations - -from typing import Final - -import pytest - -from shared.services.retrieval.nav.nav_map_scores import _pool_unit_scores_to_tree - - -_CASES: Final[tuple[tuple[dict[str, list[str]], set[str], dict[str, float]], ...]] = ( - ( - {"root-a": ["section-a", "section-b"], "section-a": [], "section-b": []}, - {"section-a", "section-b"}, - {"section-a": 0.4, "section-b": 0.8, "root-a__self": 0.2}, - ), - ( - { - "root-a": ["parent-a"], - "parent-a": ["leaf-a", "leaf-b"], - "leaf-a": [], - "leaf-b": [], - "root-b": ["leaf-c"], - "leaf-c": [], - }, - {"leaf-a", "leaf-b", "leaf-c"}, - { - "leaf-a": 0.9, - "leaf-b": 0.3, - "leaf-c": 0.7, - "parent-a__self": 0.95, - }, - ), -) - - -def _legacy_pool( - children_map: dict[str, list[str]], - leaves: set[str], - unit_scores: dict[str, float], -) -> dict[str, float]: - map_scores = { - leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves - } - - def score_node(section_id: str) -> float: - if section_id in map_scores: - return map_scores[section_id] - children = children_map.get(section_id) or [] - if not children: - score = float(unit_scores.get(section_id, 0.0) or 0.0) - map_scores[section_id] = score - return score - descendants: list[str] = [] - - def collect(section: str) -> None: - nested = children_map.get(section) or [] - if not nested: - if section in leaves: - descendants.append(section) - return - for child in nested: - collect(child) - - collect(section_id) - parts = [float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendants] - self_key = f"{section_id}__self" - if self_key in unit_scores: - parts.append(float(unit_scores[self_key])) - score = float(max(parts)) if parts else 0.0 - map_scores[section_id] = score - return score - - for section_id in children_map: - score_node(section_id) - return map_scores - - -@pytest.mark.parametrize("children_map, leaves, unit_scores", _CASES) -def test_map_score_pooling_preserves_legacy_semantics( - children_map: dict[str, list[str]], - leaves: set[str], - unit_scores: dict[str, float], -) -> None: - assert _pool_unit_scores_to_tree(children_map, leaves, unit_scores) == _legacy_pool( - children_map, leaves, unit_scores - ) diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_map_unit_index_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_map_unit_index_contract.py deleted file mode 100644 index 9cde1b65e..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_map_unit_index_contract.py +++ /dev/null @@ -1,783 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager -from typing import Any -from uuid import uuid4 - -from httpx import AsyncClient -from sqlalchemy import Engine, delete, event, select, text - -from shared.models.database.document import ( - DocumentMapUnit, - DocumentMapUnitIndex, - DocumentMapUnitToken, -) -from shared.services.retrieval.nav._compat import Chunk, EpisodeResult -from shared.services.retrieval.nav import nav_knowhere -from shared.services.retrieval.nav.nav_knowhere import ( - LazyKnowhereProvider, - NamespaceKnowhereProvider, - ReadOnlyChunkStore, -) -from shared.services.retrieval.nav.nav_map_scores import ( - compute_corpus_map_and_unit_scores, - select_map_highlights, -) -from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows -from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, - UnitRow, -) -from shared.services.retrieval.scoring.score_units import build_score_units -from shared.services.retrieval.nav_snapshot import load_nav_snapshot -from shared.services.retrieval.publication_content import ( - replace_document_revision_content, -) -from shared.services.retrieval.nav_bridge import build_referenced_chunks -from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.retrieval.serving_generation import lock_namespace_generation -from tests.support.contract_database import ContractDatabase -from tests.support.retrieval_snapshot_support import contract_db_session - -_USER_ID = "local-dev-user" - - -def test_read_only_score_loader_drives_frequency_lookup_from_token_hash( - monkeypatch, -) -> None: - document_id = "doc-frequency" - job_result_id = "revision-frequency" - executions: list[tuple[str, object]] = [] - - class FakeCursor: - def __init__(self) -> None: - self.rows: list[tuple[object, ...]] = [] - - def execute(self, statement: str, parameters: object = None) -> None: - executions.append((statement, parameters)) - if "document_map_unit_indexes" in statement: - self.rows = [ - (document_id, job_result_id, 2, 1, 0.0, 0.0, 1, 1, 1, 1) - ] - elif "FROM document_sections" in statement: - self.rows = [(document_id, job_result_id, 1)] - elif "FROM document_map_units AS units" in statement: - self.rows = [("unit-frequency", document_id, "chunk-frequency", "section-frequency", 1, 1)] - elif "FROM document_map_unit_tokens" in statement: - self.rows = [("unit-frequency", "path", "retrieval", 1)] - else: - self.rows = [] - - def fetchall(self) -> list[tuple[object, ...]]: - return self.rows - - def close(self) -> None: - return None - - class FakeConnection: - def __init__(self) -> None: - self.cursor_instance = FakeCursor() - - def set_session(self, *, readonly: bool, autocommit: bool) -> None: - assert readonly is True - assert autocommit is True - - def cursor(self) -> FakeCursor: - return self.cursor_instance - - def close(self) -> None: - return None - - connection = FakeConnection() - monkeypatch.setattr(nav_knowhere, "_connect", lambda _dsn: connection) - store = nav_knowhere.ReadOnlyChunkStore( - dsn="postgresql://test", - revisions={document_id: job_result_id}, - ) - - corpus = store.load_persisted_score_corpus( - [document_id], - {document_id: ["section-frequency"]}, - ["retrieval"], - ) - - assert corpus is not None - selective_executions = [ - statement - for statement, _parameters in executions - if "SELECT DISTINCT map_unit_id" in statement - ] - assert len(selective_executions) == 1 - assert "JOIN (VALUES" in selective_executions[0] - frequency_executions = [ - (statement, parameters) - for statement, parameters in executions - if "FROM document_map_unit_tokens" in statement - and "SELECT DISTINCT map_unit_id" not in statement - ] - assert len(frequency_executions) == 1 - statement, parameters = frequency_executions[0] - assert "scoped_units AS MATERIALIZED" in statement - assert "JOIN scoped_units" in statement - assert "channel = ANY" in statement - assert "token_hash = ANY" in statement - assert "map_unit_id = ANY" not in statement - assert isinstance(parameters, list) - assert parameters[0] == ["unit-frequency"] - assert parameters[1] == ["path", "content"] - assert parameters[2] == [ - "6e51d6a3d90b6a3243d38e6da6b3f31f49867c1360beba83da8ca9630f9672c7" - ] - - -class _IncompleteIndexStore: - """Minimal lazy store whose missing index returns no persisted scores.""" - - def __init__(self, units_by_section: Mapping[str, Sequence[UnitRow]]) -> None: - self.units_by_section = { - str(section_id): list(units) - for section_id, units in units_by_section.items() - } - self.persisted_loads = 0 - - def load_persisted_score_corpus( - self, - document_ids: Sequence[str], - allowed_section_ids_by_document: Mapping[str, Sequence[str]], - queries: Sequence[str], - ) -> None: - del document_ids, allowed_section_ids_by_document, queries - self.persisted_loads += 1 - return None - - def load_section_units( - self, - document_id: str, - section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> list[UnitRow]: - del document_id, extra_chunk_ids - return list(self.units_by_section.get(str(section_id), ())) - - def close(self) -> None: - return None - - -async def test_published_map_units_preserve_scores_without_chunk_payload_reads( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch, -) -> None: - identifier = uuid4().hex[:8] - namespace = f"map-unit-index-{identifier}" - document_id = f"doc_map_{identifier}" - job_id = f"job_map_{identifier}" - job_result_id = f"result_map_{identifier}" - async with developer_api_client_factory(): - await _seed_revision( - namespace=namespace, - document_id=document_id, - job_id=job_id, - job_result_id=job_result_id, - ) - scope = DocumentPublicationScope( - user_id=_USER_ID, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - source_file_name="indexed.pdf", - ) - chunks = [ - { - "chunk_id": "parent-a", - "type": "text", - "content": "common alpha parent evidence", - "path": "indexed.pdf/Root/Parent/intro-a", - "order": 1, - "metadata": {}, - }, - { - "chunk_id": "parent-b", - "type": "text", - "content": "common beta parent evidence", - "path": "indexed.pdf/Root/Parent/intro-b", - "order": 2, - "metadata": {}, - }, - { - "chunk_id": "leaf-a", - "type": "text", - "content": "common alpha leaf evidence", - "path": "indexed.pdf/Root/Parent/Leaf A/body", - "order": 3, - "metadata": {}, - }, - { - "chunk_id": "leaf-b", - "type": "text", - "content": "common beta leaf evidence", - "path": "indexed.pdf/Root/Parent/Leaf B/body", - "order": 4, - "metadata": {}, - }, - { - "chunk_id": "leaf-c", - "type": "text", - "content": "gamma other evidence", - "path": "indexed.pdf/Root/Parent/Leaf C/body", - "order": 5, - "metadata": {}, - }, - ] - async with contract_db_session() as db: - await db.run_sync( - lambda sync_db: _publish_revision_with_generation_lock( - sync_db, - scope=scope, - chunks=chunks, - ) - ) - await db.commit() - - async with contract_db_session() as db: - eager_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - ) - eager_toolspace = ProviderToolSpace(eager_snapshot.provider) - expected_units = build_score_units(eager_toolspace, document_id) - - async with contract_db_session() as db: - index = ( - await db.execute( - select(DocumentMapUnitIndex).where( - DocumentMapUnitIndex.document_id == document_id - ) - ) - ).scalar_one() - persisted_units = list( - ( - await db.execute( - select(DocumentMapUnit) - .where(DocumentMapUnit.document_id == document_id) - .order_by(DocumentMapUnit.sort_order) - ) - ).scalars() - ) - persisted_tokens = list( - ( - await db.execute( - select(DocumentMapUnitToken).where( - DocumentMapUnitToken.map_unit_id.in_( - [unit.id for unit in persisted_units] - ) - ) - ) - ).scalars() - ) - index_names = { - str(row[0]) - for row in ( - await db.execute( - text( - "SELECT indexname FROM pg_indexes " - "WHERE schemaname = current_schema() " - "AND tablename = 'document_map_unit_tokens'" - ) - ) - ).all() - } - lazy_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - lazy=True, - ) - - assert index.unit_count == len(expected_units) - assert index.path_document_count == sum( - unit.path_token_count > 0 for unit in persisted_units - ) - assert index.path_total_length == sum( - unit.path_token_count for unit in persisted_units - ) - assert index.content_document_count == sum( - unit.content_token_count > 0 for unit in persisted_units - ) - assert index.content_total_length == sum( - unit.content_token_count for unit in persisted_units - ) - assert [unit.unit_id for unit in persisted_units] == [ - str(unit["chunk_id"]) for unit in expected_units - ] - assert persisted_tokens - assert "idx_document_map_unit_tokens_lookup" in index_names - - def reject_payload_read( - _store: ReadOnlyChunkStore, - _document_id: str, - _section_id: str, - extra_chunk_ids: Sequence[str] = (), - ) -> list[UnitRow]: - del extra_chunk_ids - raise AssertionError("persisted map scoring loaded full chunk payloads") - - original_payload_loader = ReadOnlyChunkStore.load_section_units - monkeypatch.setattr( - ReadOnlyChunkStore, - "load_section_units", - reject_payload_read, - ) - actual_scores = compute_corpus_map_and_unit_scores( - ProviderToolSpace(lazy_snapshot.provider), - doc_ids=[document_id], - query="common alpha", - ) - monkeypatch.setattr( - ReadOnlyChunkStore, - "load_section_units", - original_payload_loader, - ) - async with contract_db_session() as db: - token_id = ( - select(DocumentMapUnitToken.id) - .join( - DocumentMapUnit, - DocumentMapUnit.id == DocumentMapUnitToken.map_unit_id, - ) - .where(DocumentMapUnit.document_id == document_id) - .limit(1) - .scalar_subquery() - ) - await db.execute( - delete(DocumentMapUnitToken).where(DocumentMapUnitToken.id == token_id) - ) - await db.execute( - delete(DocumentMapUnitIndex).where( - DocumentMapUnitIndex.document_id == document_id - ) - ) - await db.commit() - incomplete_snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - lazy=True, - ) - fallback_scores = compute_corpus_map_and_unit_scores( - ProviderToolSpace(incomplete_snapshot.provider), - doc_ids=[document_id], - query="common alpha", - ) - incomplete_snapshot.close() - lazy_snapshot.close() - eager_snapshot.close() - - assert any(score > 0.0 for score in actual_scores[1].values()) - assert select_map_highlights(actual_scores[1], k=3) - assert any(score > 0.0 for score in fallback_scores[1].values()) - assert any(score > 0.0 for score in fallback_scores[0].values()) - - -async def test_lazy_snapshot_defers_selected_asset_reference_metadata( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch, -) -> None: - identifier = uuid4().hex[:8] - namespace = f"lazy-ref-{identifier}" - document_id = f"doc_ref_{identifier}" - job_id = f"job_ref_{identifier}" - job_result_id = f"result_ref_{identifier}" - async with developer_api_client_factory(): - await _seed_revision( - namespace=namespace, - document_id=document_id, - job_id=job_id, - job_result_id=job_result_id, - ) - scope = DocumentPublicationScope( - user_id=_USER_ID, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - source_file_name="refs.pdf", - ) - chunks = [ - { - "chunk_id": "body", - "type": "text", - "content": "body evidence", - "path": "refs.pdf/Root/Section/body", - "order": 1, - "metadata": {"connect_to": [{"target": "asset"}]}, - }, - { - "chunk_id": "asset", - "type": "image", - "content": "image description", - "path": "refs.pdf/Root/image", - "order": 2, - "file_path": "images/asset.png", - "metadata": {}, - }, - ] - async with contract_db_session() as db: - await db.run_sync( - lambda sync_db: _publish_revision_with_generation_lock( - sync_db, - scope=scope, - chunks=chunks, - ) - ) - await db.commit() - - async with contract_db_session() as db: - persisted_units = list( - ( - await db.execute( - select(DocumentMapUnit).where( - DocumentMapUnit.document_id == document_id - ) - ) - ).scalars() - ) - assert any(unit.has_image for unit in persisted_units) - assert all(not unit.has_table for unit in persisted_units) - - calls: list[tuple[str, str]] = [] - original = ReadOnlyChunkStore.load_chunk_reference_metadata - - def record_reference_load( - store: ReadOnlyChunkStore, - document: str, - chunk: str, - ) -> Mapping[str, object] | None: - calls.append((document, chunk)) - return original(store, document, chunk) - - monkeypatch.setattr( - ReadOnlyChunkStore, - "load_chunk_reference_metadata", - record_reference_load, - ) - async with contract_db_session() as db: - snapshot = await load_nav_snapshot( - db, - user_id=_USER_ID, - namespace=namespace, - lazy=True, - ) - - assert calls == [] - assert snapshot.chunk_ref_index[f"{document_id}:asset"]["file_path"] == ( - "images/asset.png" - ) - assert calls == [(document_id, "asset")] - episode = EpisodeResult( - representation="", - steps=[], - scored_chunks=[ - ( - Chunk( - node_id="asset", - doc_id=document_id, - text="image description", - line_ids=(2,), - section_id="root", - ), - 0.75, - ) - ], - kept_chunks=[ - Chunk( - node_id="asset", - doc_id=document_id, - text="image description", - line_ids=(2,), - section_id="root", - ) - ], - evidence_text="image description", - evidence_chars_actual=17, - retrieved_nodes=["asset"], - ) - refs, scores = build_referenced_chunks(episode, snapshot) - assert refs == [ - { - "chunk_id": "asset", - "document_id": document_id, - "chunk_type": "image", - "section_path": "Root / image", - "file_path": "images/asset.png", - "job_id": job_id, - "score": 0.75, - } - ] - assert scores == {"asset": 0.75} - snapshot.close() - - -async def test_connected_hydration_does_not_load_legacy_job_chunks( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - identifier = uuid4().hex[:8] - namespace = f"connected-job-{identifier}" - document_id = f"doc_connected_{identifier}" - job_id = f"job_connected_{identifier}" - job_result_id = f"result_connected_{identifier}" - statements: list[str] = [] - - def capture_job_chunk_query( - _connection: Any, - _cursor: Any, - statement: str, - _parameters: Any, - _context: Any, - _executemany: bool, - ) -> None: - if "job_chunks" in statement.lower(): - statements.append(statement) - - async with developer_api_client_factory(): - await _seed_revision( - namespace=namespace, - document_id=document_id, - job_id=job_id, - job_result_id=job_result_id, - ) - scope = DocumentPublicationScope( - user_id=_USER_ID, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - source_file_name="connected.pdf", - ) - chunks = [ - { - "chunk_id": "body-connected", - "type": "text", - "content": "body connected evidence", - "path": "connected.pdf/Root/Section/body", - "order": 1, - "metadata": {"connect_to": [{"target": "asset-connected"}]}, - }, - { - "chunk_id": "asset-connected", - "type": "image", - "content": "asset connected summary", - "path": "images/asset-connected.png", - "order": 2, - "file_path": "images/asset-connected.png", - "metadata": {}, - }, - ] - async with contract_db_session() as db: - await db.run_sync( - lambda sync_db: _publish_revision_with_generation_lock( - sync_db, - scope=scope, - chunks=chunks, - ) - ) - await db.commit() - - event.listen(Engine, "before_cursor_execute", capture_job_chunk_query) - try: - async with contract_db_session() as db: - hydrated = await hydrate_connected_target_rows( - db=db, - rows=[ - { - "document_id": document_id, - "job_result_id": job_result_id, - "chunk_id": "body-connected", - "chunk_type": "text", - "chunk_metadata": { - "connect_to": [{"target": "asset-connected"}] - }, - } - ], - exclude_document_ids=[], - exclude_sections=[], - revision_pins={document_id: job_result_id}, - ) - finally: - event.remove(Engine, "before_cursor_execute", capture_job_chunk_query) - - assert [row["chunk_id"] for row in hydrated] == ["asset-connected"] - assert hydrated[0]["job_id"] == job_id - assert statements == [] - - -def test_incomplete_index_returns_empty_scores() -> None: - first_sections = [ - SectionRow("root-a", None, "Root A", "Root A", 0, "", 0), - SectionRow("leaf-a", "root-a", "Root A / Leaf A", "Leaf A", 1, "", 1), - ] - second_sections = [ - SectionRow("root-b", None, "Root B", "Root B", 0, "", 0), - SectionRow("leaf-b", "root-b", "Root B / Leaf B", "Leaf B", 1, "", 1), - ] - first_unit = UnitRow("same-chunk", "leaf-a", "text", "alpha evidence", 1) - second_unit = UnitRow("same-chunk", "leaf-b", "text", "beta evidence", 1) - - eager = ProviderToolSpace( - NamespaceKnowhereProvider( - [ - KnowhereProvider( - doc_id="doc-a", sections=first_sections, units=[first_unit] - ), - KnowhereProvider( - doc_id="doc-b", sections=second_sections, units=[second_unit] - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - ) - ) - store = _IncompleteIndexStore( - {"leaf-a": [first_unit], "leaf-b": [second_unit]} - ) - lazy = ProviderToolSpace( - NamespaceKnowhereProvider( - [ - LazyKnowhereProvider( - doc_id="doc-a", - sections=first_sections, - chunk_store=store, - known_chunk_ids=[first_unit.chunk_id], - ), - LazyKnowhereProvider( - doc_id="doc-b", - sections=second_sections, - chunk_store=store, - known_chunk_ids=[second_unit.chunk_id], - ), - ], - titles={"doc-a": "Document A", "doc-b": "Document B"}, - chunk_owner_by_id={"same-chunk": "doc-a"}, - ) - ) - - expected = compute_corpus_map_and_unit_scores( - eager, doc_ids=["doc-a", "doc-b"], query="alpha beta" - ) - actual = compute_corpus_map_and_unit_scores( - lazy, doc_ids=["doc-a", "doc-b"], query="alpha beta" - ) - - assert set(actual[1]) == {"leaf-a", "leaf-b"} - assert set(expected[1]) == {"leaf-a", "leaf-b"} - assert all(score == 0.0 for score in actual[0].values()) - assert store.persisted_loads == 1 - - -def test_titleless_leaf_has_identical_eager_and_lazy_path_scoring() -> None: - sections = [ - SectionRow("root", None, "Root", "Root", 0, "", 0), - SectionRow("leaf", "root", "Root / Leaf", "", 1, "", 1), - ] - unit = UnitRow("titleless-chunk", "leaf", "text", "alpha evidence", 1) - eager = ProviderToolSpace( - KnowhereProvider(doc_id="doc", sections=sections, units=[unit]) - ) - store = _IncompleteIndexStore({"leaf": [unit]}) - lazy = ProviderToolSpace( - LazyKnowhereProvider( - doc_id="doc", - sections=sections, - chunk_store=store, - known_chunk_ids=[unit.chunk_id], - ) - ) - - assert build_score_units(eager, "doc") == build_score_units(lazy, "doc") - assert compute_corpus_map_and_unit_scores( - eager, doc_ids=["doc"], query="alpha" - ) == compute_corpus_map_and_unit_scores( - lazy, doc_ids=["doc"], query="alpha" - ) - - -def _publish_revision_with_generation_lock( - sync_db: Any, - *, - scope: DocumentPublicationScope, - chunks: list[dict[str, Any]], -) -> None: - """Mirror the production publish sequence: lock, then patch the snapshot. - - ``replace_document_revision_content`` requires the namespace generation - row to already exist (``publication_service.py`` locks it before every - real publish); this test helper reproduces that precondition. - """ - lock_namespace_generation( - sync_db, user_id=scope.user_id, namespace=scope.namespace - ) - replace_document_revision_content(sync_db, scope=scope, chunks=chunks) - - -async def _seed_revision( - *, - namespace: str, - document_id: str, - job_id: str, - job_result_id: 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 - ) VALUES ( - :job_id, :user_id, 'document_ingestion', 'done', 'file', 0, - false, NOW(), NOW(), 0, 'skipped' - ) - """, - {"job_id": job_id, "user_id": _USER_ID}, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, source_file_name, - parse_track, created_at, updated_at - ) VALUES ( - :document_id, :user_id, :namespace, 'active', - 'indexed.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 ( - :job_result_id, :job_id, :document_id, 'inline', NOW(), NOW() - ) - """, - { - "job_result_id": job_result_id, - "job_id": job_id, - "document_id": document_id, - }, - ) - await ContractDatabase.execute( - """ - UPDATE documents SET current_job_result_id = :job_result_id - WHERE document_id = :document_id - """, - {"job_result_id": job_result_id, "document_id": document_id}, - ) diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py deleted file mode 100644 index bf72b95a9..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py +++ /dev/null @@ -1,226 +0,0 @@ -"""Map-nav route session lifetime: route rollback before fresh final hydration.""" - -from __future__ import annotations - -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import Any, cast - -import pytest -from pytest import MonkeyPatch -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.execution import routes as route_module -from shared.services.retrieval.execution.reference_resolver import ( - ResolvedWorkflowReferences, -) -from shared.services.retrieval.execution.route_types import RetrievalRouteContext -from shared.services.retrieval.nav._compat import AgentStep, Chunk, EpisodeResult -from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow -from shared.services.retrieval.nav_snapshot import build_nav_snapshot - -RouteRow = dict[str, object] - - -class _RecordingRouteSession: - def __init__(self, events: list[str]) -> None: - self._events = events - - async def rollback(self) -> None: - self._events.append("route_rollback") - - -@pytest.mark.asyncio -async def test_mapnav_route_should_release_route_session_before_fresh_final_hydration( - monkeypatch: MonkeyPatch, -) -> None: - events: list[str] = [] - route_db = _RecordingRouteSession(events) - fresh_db = object() - - snap = build_nav_snapshot( - document_titles={"doc_contract": "Contract"}, - sections_by_doc={ - "doc_contract": [ - SectionRow( - section_id="sec_contract", - parent_section_id=None, - section_path="contract/section", - section_title="section", - section_level=0, - summary="", - sort_order=0, - ) - ] - }, - units_by_doc={ - "doc_contract": [ - UnitRow( - chunk_id="chunk_contract", - section_id="sec_contract", - chunk_type="text", - content="contract content", - sort_order=0, - ) - ] - }, - chunk_ref_index={ - "chunk_contract": { - "document_id": "doc_contract", - "section_path": "contract/section", - "chunk_type": "text", - "file_path": None, - "job_id": "job_contract", - } - }, - ) - - @asynccontextmanager - async def fake_get_db_context() -> AsyncGenerator[AsyncSession, None]: - events.append("fresh_db_open") - try: - yield cast(AsyncSession, fresh_db) - finally: - events.append("fresh_db_close") - - async def fake_load_nav_snapshot( - db: AsyncSession, - **_kwargs: Any, - ) -> Any: - assert db is route_db - events.append("snapshot_loaded") - return snap - - def fake_run_nav_episode(*_args: Any, **_kwargs: Any) -> EpisodeResult: - events.append("nav_episode") - assert events[:3] == [ - "snapshot_loaded", - "route_rollback", - "nav_episode", - ] - chunk = Chunk( - node_id="chunk_contract", - doc_id="doc_contract", - text="contract content", - line_ids=(0,), - section_id="sec_contract", - ) - return EpisodeResult( - representation="mapnav", - steps=[ - AgentStep( - step_idx=1, - action="query_plan", - detail={ - "plan": {"subgoals": [], "coverage_checklist": []}, - "token_limit": 100000, - "tokens_used_total": 1, - "tokens_used_delta": 1, - "elapsed_ms": 1, - }, - ) - ], - scored_chunks=[(chunk, 1.0)], - kept_chunks=[chunk], - evidence_text="contract content", - evidence_chars_actual=len("contract content"), - retrieved_nodes=["chunk_contract"], - stop_reason="completed", - ) - - async def fake_resolve_workflow_references( - *, - db: AsyncSession, - user_id: str, - namespace: str, - refs: list[RouteRow], - score_by_chunk_id: dict[str, float] | None = None, - revision_pins: dict[str, str] | None = None, - ) -> ResolvedWorkflowReferences: - assert db is fresh_db - assert user_id == "contract-user" - assert namespace == "contract-namespace" - assert refs - assert score_by_chunk_id is not None - assert revision_pins is None - events.append("resolve_references") - row = { - "document_id": "doc_contract", - "chunk_id": "chunk_contract", - "source_file_name": "contract.pdf", - "section_path": "contract/section", - "chunk_type": "text", - "content": "contract content", - } - return ResolvedWorkflowReferences(refs=refs, rows=[row]) - - async def fake_assemble_retrieval_results( - *, - db: AsyncSession, - rows: list[RouteRow], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - allowed_chunk_types: set[str] | None, - revision_pins: dict[str, str] | None = None, - ) -> list[RouteRow]: - assert db is fresh_db - assert exclude_document_ids == [] - assert exclude_sections == [] - assert allowed_chunk_types is None - assert revision_pins is None - events.append("assemble_results") - return rows - - monkeypatch.setattr( - "shared.services.retrieval.nav_snapshot.load_nav_snapshot", - fake_load_nav_snapshot, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.run_nav_episode", - fake_run_nav_episode, - ) - monkeypatch.setattr(route_module, "open_fresh_database_context", fake_get_db_context) - monkeypatch.setattr( - route_module, - "resolve_workflow_references", - fake_resolve_workflow_references, - ) - monkeypatch.setattr( - route_module, - "assemble_retrieval_results", - fake_assemble_retrieval_results, - ) - - outcome = await route_module._run_mapnav_route( - RetrievalRouteContext( - db=cast(AsyncSession, route_db), - user_id="contract-user", - namespace="contract-namespace", - query="session lifetime", - top_k=1, - exclude_document_ids=[], - exclude_sections=[], - allowed_chunk_types=None, - chunk_types=None, - signal_paths=None, - filter_mode="delete", - rerank=False, - threshold=0.0, - internal_recall_k=None, - effective_recall_k=3, - use_agentic=True, - ) - ) - - assert events == [ - "snapshot_loaded", - "route_rollback", - "nav_episode", - "fresh_db_open", - "resolve_references", - "assemble_results", - "fresh_db_close", - ] - assert outcome.response["router_used"] == "mapnav" - assert outcome.response["results"][0]["document_id"] == "doc_contract" - assert outcome.completion_label == "MAPNAV RETRIEVAL" diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_relit_map_cache_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_relit_map_cache_contract.py deleted file mode 100644 index b8257c9c0..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_relit_map_cache_contract.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -from shared.services.retrieval.nav.nav_orchestrate import _relit_map -from shared.services.retrieval.nav.nav_types import NavState - - -def test_relit_map_reuses_same_query_within_episode(monkeypatch) -> None: - calls: list[str] = [] - - def fake_relight_map_for_query(*_args, **kwargs): - calls.append(str(kwargs["query"])) - return {"section": 1.0}, {"unit": 2.0}, ["section"] - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_map_scores.relight_map_for_query", - fake_relight_map_for_query, - ) - state = NavState(doc_id="", query="retrieval") - config = type("Config", (), {"collect_top_k": 6})() - - with _relit_map(None, state, config, query="retrieval"): - pass - with _relit_map(None, state, config, query="retrieval"): - pass - - assert calls == ["retrieval"] - assert state.relit_map_cache["retrieval"][0] == {"section": 1.0} diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_batching_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_batching_contract.py deleted file mode 100644 index a7239100a..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_batching_contract.py +++ /dev/null @@ -1,194 +0,0 @@ -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/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_consistency_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_consistency_contract.py deleted file mode 100644 index a30af3c31..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_consistency_contract.py +++ /dev/null @@ -1,264 +0,0 @@ -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 -import pytest -from sqlalchemy import Executable, Result -from sqlalchemy.exc import SQLAlchemyError - -from shared.services.retrieval.execution.reference_resolver import ( - resolve_workflow_references, -) -from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot -from shared.services.retrieval.nav_snapshot import _resolve_namespace_snapshot_entries -from tests.support.retrieval_snapshot_support import contract_db_session -from tests.support.contract_database import ContractDatabase - - -_USER_ID = "local-dev-user" - - -class _GenerationUnavailableSession: - def __init__(self) -> None: - self.rollback_count = 0 - - async def execute(self, _statement: Executable) -> Result[tuple[object, ...]]: - raise SQLAlchemyError("generation table unavailable") - - async def rollback(self) -> None: - self.rollback_count += 1 - - -@pytest.mark.asyncio -async def test_snapshot_loader_falls_back_when_generation_cannot_be_verified() -> None: - session = _GenerationUnavailableSession() - - result = await _resolve_namespace_snapshot_entries( - session, - user_id=_USER_ID, - namespace="default", - document_revisions=[("doc-a", "result-a")], - ) - - assert result is None - assert session.rollback_count == 1 - - -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" - ) - - -@pytest.mark.asyncio -async def test_reference_hydration_keeps_the_captured_revision_after_republish( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - namespace = f"revision-hydration-race-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - document_id, new_result_id = await _seed_republished_document(namespace) - revision_rows = await ContractDatabase.fetch_all( - """ - SELECT job_result_id, chunk_id - FROM document_chunks - WHERE document_id = :document_id - ORDER BY job_result_id - """, - {"document_id": document_id}, - ) - old_revision = next( - row for row in revision_rows if row["job_result_id"] != new_result_id - ) - 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: - resolved = await resolve_workflow_references( - db=db, - user_id=_USER_ID, - namespace=namespace, - refs=[ - { - "document_id": document_id, - "chunk_id": old_revision["chunk_id"], - } - ], - revision_pins={document_id: old_revision["job_result_id"]}, - ) - - assert [row["content"] for row in resolved.rows] == ["old content"] - assert [row["job_result_id"] for row in resolved.rows] == [ - old_revision["job_result_id"] - ] diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_large_corpus_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_large_corpus_contract.py deleted file mode 100644 index f200ff845..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_large_corpus_contract.py +++ /dev/null @@ -1,358 +0,0 @@ -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 -import pytest - -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, - _CHUNK_BATCH_SIZE, - _REVISION_GROUP_SIZE, - load_nav_snapshot, -) -from shared.services.retrieval import nav_snapshot as nav_snapshot_module -from sqlalchemy import Executable, Result, select -from sqlalchemy.engine import Row -from sqlalchemy.sql.selectable import Select -from tests.support.retrieval_snapshot_support import contract_db_session -from tests.support.contract_database import ContractDatabase - - -_USER_ID = "local-dev-user" -_DOCUMENT_COUNT = 100 -_CHUNKS_PER_DOCUMENT = 600 -_SECTIONS_PER_DOCUMENT = 8 -_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, - ] -] - - -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, ...]]: - 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) - - async def rollback(self) -> None: - await self._session.rollback() - - -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, -) -> 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: - rows: list[LegacySnapshotRow] = list((await db.execute(stmt)).all()) - return rows - - -async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: pytest.MonkeyPatch, -) -> None: - namespace = f"large-corpus-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - await _seed_large_retrieval_corpus(namespace) - async def unexpected_manifest_load(*_args: object, **_kwargs: object) -> None: - raise AssertionError("large snapshots must use normalized retrieval rows") - - monkeypatch.setattr( - nav_snapshot_module, - "_load_manifest_sections", - unexpected_manifest_load, - ) - legacy_rows = await _load_legacy_rows(namespace) - async with contract_db_session() as db: - counting_db = _CountingSession(db) - snapshot = await load_nav_snapshot( - counting_db, - user_id=_USER_ID, - namespace=namespace, - ) - 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 - - async with contract_db_session() as db: - 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_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) - ] - 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, - ) - 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 diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_redis_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_redis_contract.py deleted file mode 100644 index 299341336..000000000 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_redis_contract.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Contract tests for binary namespace snapshot caching.""" - -from __future__ import annotations - -import pytest -from sqlalchemy import Executable - -from shared.core.config.redis import RedisConfig, RedisConfigManager -from shared.services.redis.redis_service_factory import RedisServiceFactory -from shared.services.redis.redis_service import RedisService -from shared.services.retrieval.nav_snapshot import _resolve_namespace_snapshot_entries -from shared.services.retrieval.namespace_map_snapshot_redis import ( - NamespaceMapSnapshotRedisCache, -) -from shared.services.retrieval.serving_manifest import encode_namespace_map_snapshot - - -class _FakeRedisClient: - def __init__(self, *, decode_responses: bool) -> None: - self.decode_responses = decode_responses - self.values: dict[str, bytes] = {} - self.ttls: dict[str, int] = {} - - async def get(self, key: str) -> bytes | None: - return self.values.get(key) - - async def set(self, key: str, value: bytes, *, ex: int) -> bool: - self.values[key] = value - self.ttls[key] = ex - return True - - async def aclose(self) -> None: - return None - - -@pytest.mark.asyncio -async def test_binary_redis_operations_preserve_compressed_snapshot_bytes( - monkeypatch: pytest.MonkeyPatch, -) -> None: - clients: list[_FakeRedisClient] = [] - - def create_client(*args: object, **kwargs: object) -> _FakeRedisClient: - client = _FakeRedisClient( - decode_responses=bool(kwargs.get("decode_responses")) - ) - clients.append(client) - return client - - monkeypatch.setattr( - "shared.services.redis.redis_service.redis.from_url", create_client - ) - service = RedisService(RedisConfigManager(RedisConfig())) - payload = b"\x78\x9c\x00\xffcompressed-snapshot" - - assert await service.set_bytes("contract:snapshot", payload, ex=3600) - assert await service.get_bytes("contract:snapshot") == payload - assert len(clients) == 1 - assert clients[0].decode_responses is False - assert clients[0].ttls["knowhere-api:contract:snapshot"] == 3600 - - await service.close() - - -@pytest.mark.asyncio -async def test_snapshot_cache_scopes_reads_and_writes_by_generation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class _FakeSnapshotService: - def __init__(self) -> None: - self.get_keys: list[str] = [] - self.set_calls: list[tuple[str, bytes, int]] = [] - - async def get_bytes(self, key: str) -> bytes | None: - self.get_keys.append(key) - return b"snapshot" - - async def set_bytes(self, key: str, value: bytes, *, ex: int) -> bool: - self.set_calls.append((key, value, ex)) - return True - - fake_service = _FakeSnapshotService() - monkeypatch.setattr( - RedisServiceFactory, - "get_service", - classmethod(lambda cls: fake_service), - ) - - assert ( - await NamespaceMapSnapshotRedisCache.get( - user_id="user", - namespace=" ", - generation=7, - ) - == b"snapshot" - ) - assert await NamespaceMapSnapshotRedisCache.set( - user_id="user", - namespace="default", - generation=8, - payload_zlib=b"compressed", - ) - - assert fake_service.get_keys == ["retrieval:snapshot:v2:user:default:g7"] - assert fake_service.set_calls == [ - ( - "retrieval:snapshot:v2:user:default:g8", - b"compressed", - 3600, - ) - ] - - -class _SnapshotResult: - def __init__(self, *, scalar: object = None, row: tuple[object, ...] | None = None): - self._scalar = scalar - self._row = row - - def scalar_one_or_none(self) -> object: - return self._scalar - - def first(self) -> tuple[object, ...] | None: - return self._row - - -class _SnapshotSequenceSession: - def __init__(self, results: list[_SnapshotResult]) -> None: - self._results = iter(results) - - async def execute(self, _statement: Executable) -> _SnapshotResult: - return next(self._results) - - async def rollback(self) -> None: - return None - - -@pytest.mark.asyncio -async def test_corrupt_redis_blob_falls_back_to_postgres_snapshot( - monkeypatch: pytest.MonkeyPatch, -) -> None: - payload_zlib, checksum, format_version = encode_namespace_map_snapshot( - { - "documents": { - "doc-a": { - "job_result_id": "result-a", - "job_id": "job-a", - "sections": [], - "chunks": [], - } - } - } - ) - set_calls: list[bytes] = [] - - async def get_corrupt_blob(**_: object) -> bytes: - return b"corrupt" - - async def record_repaired_blob(**kwargs: object) -> bool: - set_calls.append(bytes(kwargs["payload_zlib"])) - return True - - monkeypatch.setattr(NamespaceMapSnapshotRedisCache, "get", get_corrupt_blob) - monkeypatch.setattr(NamespaceMapSnapshotRedisCache, "set", record_repaired_blob) - session = _SnapshotSequenceSession( - [ - _SnapshotResult(scalar=3), - _SnapshotResult(row=(3, checksum, format_version)), - _SnapshotResult(row=(payload_zlib,)), - ] - ) - - entries = await _resolve_namespace_snapshot_entries( - session, - user_id="user", - namespace="default", - document_revisions=[("doc-a", "result-a")], - expected_generation=3, - ) - - assert entries is not None - assert entries[0][0:2] == ("doc-a", "result-a") - assert set_calls == [payload_zlib] diff --git a/deprecated/mapnav/tests/shared/test_nav_bridge_config.py b/deprecated/mapnav/tests/shared/test_nav_bridge_config.py deleted file mode 100644 index 6eb90319f..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_bridge_config.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Unit tests for map-nav config + exit bridge.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from typing import Any, List, Tuple - -os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") -os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") -os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") -os.environ.setdefault("S3_ACCESS_KEY_ID", "test") -os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") -os.environ.setdefault("S3_TEMP_PATH", "/tmp") - -from shared.services.retrieval.nav._compat import Chunk -from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow -from shared.services.retrieval.nav_bridge import build_referenced_chunks -from shared.services.retrieval.nav_config import ( - MAPNAV_MODEL, - build_nav_config, - nav_evidence_chars, -) -from shared.services.retrieval.nav_snapshot import build_nav_snapshot - - -def test_build_nav_config_is_checklist_map_trim_stack() -> None: - cfg = build_nav_config() - assert cfg.mode == "checklist" - assert cfg.is_checklist - assert cfg.policy == "llm" - assert cfg.subgoal_max_attempts == 2 - assert cfg.max_replans == 1 - assert cfg.max_waves == 0 - assert cfg.max_harvest_depth == 5 - assert not hasattr(cfg, "compose_packing_mode") - assert cfg.llm_model == MAPNAV_MODEL - assert cfg.planner_model == MAPNAV_MODEL - assert cfg.subagent_model == MAPNAV_MODEL - assert cfg.planner_thinking == "enabled" - assert cfg.planner_think_max_tokens == 16_384 - assert cfg.token_limit == 100_000 - assert cfg.enable_node_filter is True - assert not hasattr(cfg, "llm_model_env") - - -def test_nav_evidence_chars_is_code_constant() -> None: - assert nav_evidence_chars() == 15000 - - -def _root_remount_snapshot(): - root = SectionRow( - section_id="sec_root", - parent_section_id=None, - section_path="Root", - section_title="Root", - section_level=0, - summary="", - sort_order=0, - ) - host = SectionRow( - section_id="sec_host", - parent_section_id="sec_root", - section_path="Chapter 1", - section_title="Chapter 1", - section_level=1, - summary="", - sort_order=1, - ) - text_unit = UnitRow( - chunk_id="chk_text", - section_id="sec_host", - chunk_type="text", - content="body", - sort_order=0, - metadata={"connect_to": [{"target": "chk_img"}]}, - ) - img_unit = UnitRow( - chunk_id="chk_img", - section_id="sec_root", - chunk_type="image", - content="a.png", - sort_order=1, - file_path="a.png", - ) - return build_nav_snapshot( - document_titles={"doc_a": "Doc A"}, - sections_by_doc={"doc_a": [root, host]}, - units_by_doc={"doc_a": [text_unit, img_unit]}, - chunk_ref_index={ - "chk_text": { - "document_id": "doc_a", - "section_path": "Chapter 1", - "chunk_type": "text", - "file_path": None, - "job_id": "job_1", - }, - "chk_img": { - "document_id": "doc_a", - "section_path": "Root", - "chunk_type": "image", - "file_path": "a.png", - "job_id": "job_1", - }, - }, - ) - - -@dataclass -class _Ep: - kept_chunks: List[Any] - scored_chunks: List[Tuple[Any, float]] = field(default_factory=list) - - -def test_bridge_expands_three_node_id_forms_and_keeps_db_section_path() -> None: - """chunk_id / leaf section_id / {sid}__self all resolve via chunk_ref_index.""" - snap = _root_remount_snapshot() - - by_chunk = build_referenced_chunks( - _Ep( - kept_chunks=[ - Chunk( - node_id="chk_text", - doc_id="doc_a", - text="body", - line_ids=(0,), - section_id="sec_host", - ) - ], - scored_chunks=[ - ( - Chunk( - node_id="chk_text", - doc_id="doc_a", - text="body", - line_ids=(0,), - section_id="sec_host", - ), - 0.9, - ) - ], - ), - snap, - ) - assert by_chunk[0][0]["chunk_id"] == "chk_text" - assert by_chunk[0][0]["section_path"] == "Chapter 1" - assert by_chunk[0][0]["job_id"] == "job_1" - assert by_chunk[1].get("chk_text") == 0.9 - - by_section = build_referenced_chunks( - _Ep( - kept_chunks=[ - Chunk( - node_id="sec_host", - doc_id="doc_a", - text="host", - line_ids=(0,), - section_id="sec_host", - ) - ] - ), - snap, - ) - section_ids = {r["chunk_id"] for r in by_section[0]} - assert "chk_text" in section_ids - # After Root-remount, image lives under host for nav but ref path stays DB Root. - assert "chk_img" in section_ids - img_ref = next(r for r in by_section[0] if r["chunk_id"] == "chk_img") - assert img_ref["section_path"] == "Root" - assert img_ref["job_id"] == "job_1" - - by_self = build_referenced_chunks( - _Ep( - kept_chunks=[ - Chunk( - node_id="sec_host__self", - doc_id="doc_a", - text="intro", - line_ids=(0,), - section_id="sec_host", - ) - ] - ), - snap, - ) - self_ids = {r["chunk_id"] for r in by_self[0]} - assert self_ids == section_ids - assert all(r.get("job_id") == "job_1" for r in by_self[0]) diff --git a/deprecated/mapnav/tests/shared/test_nav_llm_backend.py b/deprecated/mapnav/tests/shared/test_nav_llm_backend.py deleted file mode 100644 index ebc3b91c4..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_llm_backend.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Unit tests for map-nav sync LLM backend (no live API).""" - -from __future__ import annotations - -import os -from typing import Any -from unittest.mock import MagicMock, 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 shared.services.retrieval import nav_llm_backend as backend -from shared.services.retrieval.nav.nav_llm import get_nav_chat_backend, nav_chat - - -def test_install_registers_backend() -> None: - assert get_nav_chat_backend() is backend.nav_chat_sync_backend - - -def test_backend_forwards_thinking_extra_body() -> None: - captured: dict[str, Any] = {} - - class FakeClient: - def chat_completion_with_usage(self, messages, **kwargs): # type: ignore[no-untyped-def] - captured["messages"] = messages - captured["kwargs"] = kwargs - return "ok", {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} - - with ( - patch( - "shared.services.ai.llm_overrides.resolve_text", - return_value=("deepseek-v4-flash", None, None), - ), - patch( - "shared.services.ai.openai_compatible_client_sync.get_openai_client", - return_value=FakeClient(), - ), - ): - out = backend.nav_chat_sync_backend( - messages=[{"role": "user", "content": "hi"}], - model="deepseek-v4-flash", - temperature=0.0, - max_tokens=256, - extra={"extra_body": {"thinking": {"type": "disabled"}}}, - timeout=30, - ) - - assert out["content"] == "ok" - assert out["usage"]["total_tokens"] == 3 - assert captured["kwargs"]["extra_body"] == {"thinking": {"type": "disabled"}} - assert captured["kwargs"]["model"] == "deepseek-v4-flash" - - -def test_nav_chat_uses_installed_backend() -> None: - fake = MagicMock( - return_value={ - "content": '{"ok": true}', - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - } - ) - from shared.services.retrieval.nav import nav_llm as nav_llm_mod - - prev = nav_llm_mod.get_nav_chat_backend() - try: - nav_llm_mod.set_nav_chat_backend(fake) - result = nav_chat( - purpose="test", - messages=[{"role": "user", "content": "x"}], - model="deepseek-v4-flash", - thinking_role="action", - ) - assert result["content"] == '{"ok": true}' - assert fake.called - call_kw = fake.call_args.kwargs - assert call_kw["extra"]["extra_body"]["thinking"]["type"] == "disabled" - finally: - nav_llm_mod.set_nav_chat_backend(prev) - - -def test_backend_uses_mapnav_model_when_model_empty() -> None: - captured: dict[str, Any] = {} - - class FakeClient: - def chat_completion_with_usage(self, messages, **kwargs): # type: ignore[no-untyped-def] - captured["kwargs"] = kwargs - return "ok", {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} - - with ( - patch( - "shared.services.ai.llm_overrides.resolve_text", - side_effect=lambda model: (model, None, None), - ), - patch( - "shared.services.ai.openai_compatible_client_sync.get_openai_client", - return_value=FakeClient(), - ), - ): - backend.nav_chat_sync_backend( - messages=[{"role": "user", "content": "hi"}], - model="", - temperature=0.0, - max_tokens=64, - ) - assert captured["kwargs"]["model"] == "deepseek-v4-flash" diff --git a/deprecated/mapnav/tests/shared/test_nav_node_filter.py b/deprecated/mapnav/tests/shared/test_nav_node_filter.py deleted file mode 100644 index bca04462c..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_node_filter.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Unit tests for deterministic WHERE node filter (in-memory hierarchy).""" - -from __future__ import annotations - -from shared.services.retrieval.nav.nav_knowhere import NamespaceKnowhereProvider -from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, -) -from shared.services.retrieval.nav.nav_node_filter import ( - apply_node_filter, - node_filter, - render_submap_observation, -) -from shared.services.retrieval.scoring.node_filter_predicates import field_predicate - - -def _section( - section_id: str, - parent: str | None, - path: str, - title: str, - *, - level: int, - summary: str = "", - order: int = 0, -) -> SectionRow: - return SectionRow( - section_id=section_id, - parent_section_id=parent, - section_path=path, - section_title=title, - section_level=level, - summary=summary, - sort_order=order, - ) - - -def _namespace_ts() -> ProviderToolSpace: - apple = KnowhereProvider( - doc_id="doc_apple", - sections=[ - _section("sec_root_a", None, "Root", "Root", level=0, order=0), - _section( - "sec_q3", - "sec_root_a", - "Q3 Results", - "Q3 Results", - level=1, - summary="Apple quarterly profit and revenue", - order=1, - ), - _section( - "sec_hw", - "sec_root_a", - "Hardware", - "Hardware", - level=1, - summary="iPhone unit sales", - order=2, - ), - ], - units=(), - ) - orange = KnowhereProvider( - doc_id="doc_orange", - sections=[ - _section("sec_root_o", None, "Root", "Root", level=0, order=0), - _section( - "sec_crop", - "sec_root_o", - "Crop Report", - "Crop Report", - level=1, - summary="orange harvest yield", - order=1, - ), - ], - units=(), - ) - provider = NamespaceKnowhereProvider( - [apple, orange], - titles={ - "doc_apple": "AAPL 10-K.pdf", - "doc_orange": "Citrus Outlook.docx", - }, - ) - return ProviderToolSpace(provider) - - -def test_path_filter_returns_complete_set_across_documents() -> None: - ts = _namespace_ts() - result = apply_node_filter( - ts, - ["doc_apple", "doc_orange"], - node_filter([field_predicate("path", ["AAPL", "Apple"])]), - ) - - assert result.failed_predicates == [] - assert result.matched_doc_ids == ["doc_apple"] - assert "sec_q3" in result.matched_section_ids - assert "sec_hw" in result.matched_section_ids - assert "sec_crop" not in result.matched_section_ids - assert result.cardinality == len(result.matched_section_ids) - assert result.cardinality >= 2 - - -def test_summary_filter_and_path_and_summary() -> None: - ts = _namespace_ts() - summary_only = apply_node_filter( - ts, - ["doc_apple", "doc_orange"], - node_filter([field_predicate("summary", ["profit"])]), - ) - assert summary_only.matched_section_ids == ["sec_q3"] - - both = apply_node_filter( - ts, - ["doc_apple", "doc_orange"], - node_filter( - [ - field_predicate("path", ["AAPL", "Apple"]), - field_predicate("summary", ["profit"]), - ] - ), - ) - assert both.matched_section_ids == ["sec_q3"] - assert both.matched_doc_ids == ["doc_apple"] - - -def test_zero_hits_and_invalid_regex_recorded() -> None: - ts = _namespace_ts() - empty = apply_node_filter( - ts, - ["doc_apple", "doc_orange"], - node_filter([field_predicate("path", ["zzz-not-present"])]), - ) - assert empty.cardinality == 0 - assert empty.matched_section_ids == [] - assert empty.matched_doc_ids == [] - - bad = apply_node_filter( - ts, - ["doc_apple", "doc_orange"], - node_filter([field_predicate("path", ["(unclosed"], match="regex")]), - ) - assert bad.cardinality == 0 - assert bad.matched_section_ids == [] - assert bad.failed_predicates == ["path:regex:invalid"] - - -def test_regex_or_terms_and_preview_budget() -> None: - ts = _namespace_ts() - result = apply_node_filter( - ts, - ["doc_apple", "doc_orange"], - node_filter( - [field_predicate("summary", ["profit|yield"], match="regex")] - ), - ) - assert set(result.matched_section_ids) == {"sec_q3", "sec_crop"} - assert result.cardinality == 2 - - preview = render_submap_observation(ts, result) - assert preview.startswith("hits=2") - assert "sec_q3" not in preview # paths shown, not ids - assert "Q3 Results" in preview diff --git a/deprecated/mapnav/tests/shared/test_nav_node_filter_wire.py b/deprecated/mapnav/tests/shared/test_nav_node_filter_wire.py deleted file mode 100644 index 1025a6473..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_node_filter_wire.py +++ /dev/null @@ -1,317 +0,0 @@ -"""Orchestrate pre-pass + allowed-scope projection tests.""" - -from __future__ import annotations - -from typing import Any - -from shared.services.retrieval.nav.nav_knowhere import NamespaceKnowhereProvider -from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, -) -from shared.services.retrieval.nav.nav_orchestrate import _execute_subgoal_harvest_once -from shared.services.retrieval.nav.nav_plan import RetrievalPlan, Subgoal -from shared.services.retrieval.nav.nav_projection import build_map -from shared.services.retrieval.nav.nav_scope_filter import ScopeFilterOutcome -from shared.services.retrieval.nav.nav_types import NavConfig, NavState - - -def _section( - section_id: str, - parent: str | None, - path: str, - title: str, - *, - level: int, - summary: str = "", - order: int = 0, -) -> SectionRow: - return SectionRow( - section_id=section_id, - parent_section_id=parent, - section_path=path, - section_title=title, - section_level=level, - summary=summary, - sort_order=order, - ) - - -def _ts() -> ProviderToolSpace: - apple = KnowhereProvider( - doc_id="doc_apple", - sections=[ - _section("sec_root_a", None, "Root", "Root", level=0, order=0), - _section( - "sec_q3", - "sec_root_a", - "Q3 Results", - "Q3 Results", - level=1, - summary="profit", - order=1, - ), - _section( - "sec_hw", - "sec_root_a", - "Hardware", - "Hardware", - level=1, - summary="iphone", - order=2, - ), - ], - units=(), - ) - return ProviderToolSpace( - NamespaceKnowhereProvider([apple], titles={"doc_apple": "AAPL 10-K.pdf"}) - ) - - -def _cfg(**kwargs: Any) -> NavConfig: - data = { - "enable_node_filter": True, - "mode": "checklist", - "llm_model": "test-model", - } - data.update(kwargs) - return NavConfig.from_dict(data) - - -def _subgoal(*, use_filter: bool = True) -> Subgoal: - return Subgoal( - id="s1", - need="apple profit", - retrieval_query="apple profit", - use_node_filter=use_filter, - ) - - -def test_collect_all_skips_harvest(monkeypatch: Any) -> None: - harvest_calls: list[dict[str, Any]] = [] - collect_calls: list[list[str]] = [] - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", - lambda *args, **kwargs: ScopeFilterOutcome( - decision="collect_all", - settled_section_ids=["sec_q3"], - settled_doc_ids=["doc_apple"], - rounds=1, - reason="small", - ), - ) - - def fake_harvest(*args: Any, **kwargs: Any) -> Any: - harvest_calls.append({"args": args, "kwargs": kwargs}) - raise AssertionError("harvest should not run") - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_harvest.harvest", - fake_harvest, - ) - - def fake_collect(ts: Any, state: NavState, chosen: Any, config: Any) -> dict[str, Any]: - del ts, config - batch = list((chosen.metadata or {}).get("batch_actions") or [chosen]) - sids = [str(a.section_id) for a in batch] - collect_calls.append(sids) - state.explicit_collect_ids.update(sids) - state.collected_section_ids.update(sids) - return {"collect_section_ids": sids} - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_navigate._apply_collect", - fake_collect, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._relit_map", - lambda *args, **kwargs: _nullcontext(), - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", - lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), - ) - - out = _execute_subgoal_harvest_once( - _ts(), - NavState(doc_id="doc_apple", query="apple profit"), - _cfg(), - RetrievalPlan(subgoals=[_subgoal()]), - _subgoal(), - steps_out=[], - ) - assert harvest_calls == [] - assert collect_calls == [["sec_q3"]] - assert out["harvest"]["reason"] == "node_filter_collect_all" - - -def test_scoped_harvest_passes_allowed_ids(monkeypatch: Any) -> None: - seen: dict[str, Any] = {} - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", - lambda *args, **kwargs: ScopeFilterOutcome( - decision="scoped_harvest", - settled_section_ids=["sec_q3"], - settled_doc_ids=["doc_apple"], - rounds=1, - reason="medium", - ), - ) - - def fake_harvest(*args: Any, **kwargs: Any) -> Any: - seen["allowed"] = kwargs.get("allowed_section_ids") - return type( - "HR", - (), - { - "n_policy_calls": 1, - "visited_section_ids": [], - "max_depth_hit": False, - "reason": "harvested", - }, - )() - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_harvest.harvest", - fake_harvest, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._relit_map", - lambda *args, **kwargs: _nullcontext(), - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", - lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), - ) - - _execute_subgoal_harvest_once( - _ts(), - NavState(doc_id="doc_apple", query="apple profit"), - _cfg(), - RetrievalPlan(subgoals=[_subgoal()]), - _subgoal(), - steps_out=[], - ) - assert seen["allowed"] == {"sec_q3"} - - -def test_fallback_matches_today(monkeypatch: Any) -> None: - seen: dict[str, Any] = {} - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", - lambda *args, **kwargs: ScopeFilterOutcome( - decision="fallback", - reason="policy_fallback", - ), - ) - - def fake_harvest(*args: Any, **kwargs: Any) -> Any: - seen["allowed"] = kwargs.get("allowed_section_ids") - return type( - "HR", - (), - { - "n_policy_calls": 1, - "visited_section_ids": [], - "max_depth_hit": False, - "reason": "", - }, - )() - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_harvest.harvest", - fake_harvest, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._relit_map", - lambda *args, **kwargs: _nullcontext(), - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", - lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), - ) - - _execute_subgoal_harvest_once( - _ts(), - NavState(doc_id="doc_apple", query="apple profit"), - _cfg(), - RetrievalPlan(subgoals=[_subgoal()]), - _subgoal(), - steps_out=[], - ) - assert seen["allowed"] is None - - -def test_flag_off_skips_pre_pass(monkeypatch: Any) -> None: - called = {"n": 0} - - def boom(*args: Any, **kwargs: Any) -> Any: - called["n"] += 1 - raise AssertionError("scope filter should not run") - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", - boom, - ) - - def fake_harvest(*args: Any, **kwargs: Any) -> Any: - return type( - "HR", - (), - { - "n_policy_calls": 1, - "visited_section_ids": [], - "max_depth_hit": False, - "reason": "", - }, - )() - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_harvest.harvest", - fake_harvest, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._relit_map", - lambda *args, **kwargs: _nullcontext(), - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", - lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), - ) - _execute_subgoal_harvest_once( - _ts(), - NavState(doc_id="doc_apple", query="apple profit"), - _cfg(), - RetrievalPlan(subgoals=[_subgoal(use_filter=False)]), - _subgoal(use_filter=False), - steps_out=[], - ) - assert called["n"] == 0 - - -def test_projection_keeps_allowed_and_ancestors() -> None: - ts = _ts() - cfg = _cfg() - projection = build_map( - ts, - doc_id="", - query="profit", - scope=None, - config=cfg, - allowed_section_ids={"sec_q3"}, - ) - ids = {v.section_id for v in projection.tree_sections} - assert "sec_q3" in ids - assert "sec_hw" not in ids - assert "sec_root_a" in ids or "doc_apple" in ids - - -class _nullcontext: - def __enter__(self) -> None: - return None - - def __exit__(self, *args: Any) -> None: - return None diff --git a/deprecated/mapnav/tests/shared/test_nav_plan_node_filter.py b/deprecated/mapnav/tests/shared/test_nav_plan_node_filter.py deleted file mode 100644 index 9099034e5..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_plan_node_filter.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Planner parse compat for use_node_filter.""" - -from __future__ import annotations - -from shared.services.retrieval.nav.nav_plan import ( - _planner_system_prompt, - parse_retrieval_plan, -) - - -def test_parse_missing_filter_flag_defaults_false() -> None: - plan = parse_retrieval_plan( - { - "reason": "old plan", - "map_coverage": "sufficient", - "subgoals": [ - { - "id": "s1", - "need": "profit", - "retrieval_query": "apple profit", - } - ], - }, - query="apple profit", - ) - assert len(plan.subgoals) == 1 - assert plan.subgoals[0].use_node_filter is False - assert plan.subgoals[0].node_filter_predicates == [] - - -def test_parse_use_node_filter_and_seed() -> None: - plan = parse_retrieval_plan( - { - "reason": "enum tickers", - "map_coverage": "sufficient", - "subgoals": [ - { - "id": "s1", - "need": "apple profit", - "retrieval_query": "AAPL profit", - "use_node_filter": True, - "node_filter": [ - { - "field": "path", - "terms": ["AAPL", "Apple"], - "match": "substring", - } - ], - } - ], - }, - query="apple profit", - ) - sg = plan.subgoals[0] - assert sg.use_node_filter is True - assert sg.node_filter_predicates == [ - {"field": "path", "terms": ["AAPL", "Apple"], "match": "substring"} - ] - - -def test_planner_prompt_mentions_where_vs_fuzzy() -> None: - text = _planner_system_prompt(max_subgoals=0) - assert "use_node_filter" in text - assert "fallback" in text - assert "only the user query" in text.lower() - assert "folded" not in text.lower() diff --git a/deprecated/mapnav/tests/shared/test_nav_plan_query_only.py b/deprecated/mapnav/tests/shared/test_nav_plan_query_only.py deleted file mode 100644 index d94ce9f24..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_plan_query_only.py +++ /dev/null @@ -1,253 +0,0 @@ -"""Planner is query-only; map lighting happens after the plan exists.""" - -from __future__ import annotations - -from typing import Any - -from shared.services.retrieval.nav.nav_orchestrate import seed_episode_map_scores_from_plan -from shared.services.retrieval.nav.nav_plan import ( - RetrievalPlan, - Subgoal, - _planner_system_prompt, - fallback_plan, -) -from shared.services.retrieval.nav.nav_types import NavConfig, NavState - - -def test_planner_system_prompt_is_query_only() -> None: - text = _planner_system_prompt(max_subgoals=3) - assert "only the user query" in text.lower() - assert "no document map" in text.lower() - assert "folded" not in text.lower() - assert "Hit" not in text - assert "collect=C" not in text - assert "whether the user query alone is enough" in text.lower() - - -def test_seed_episode_scores_plan_retrieval_queries(monkeypatch: Any) -> None: - seen: dict[str, Any] = {} - - def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): - seen["queries"] = list(queries) - seen["top_k"] = top_k - out = {} - for i, q in enumerate(queries): - out[q] = ({f"n{i}": 1.0}, {f"u{i}": 2.0}, [f"u{i}"]) - return out - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", - fake_relight, - ) - state = NavState(doc_id="", query="user original", task_type="unknown") - state.relit_map_cache = {"stale": ({}, {}, [])} - plan = RetrievalPlan( - subgoals=[ - Subgoal(id="s1", need="a", retrieval_query="alpha terms"), - Subgoal(id="s2", need="b", retrieval_query="beta terms"), - Subgoal(id="s3", need="c", retrieval_query="alpha terms"), - ] - ) - n = seed_episode_map_scores_from_plan( - object(), state, NavConfig(collect_top_k=7), plan - ) - assert n == 2 - assert seen["queries"] == ["alpha terms", "beta terms"] - assert seen["top_k"] == 7 - assert "stale" not in state.relit_map_cache - assert state.map_scores == {"n0": 1.0} - assert state.unit_scores == {"u0": 2.0} - assert state.highlight_ids == ["u0"] - assert "alpha terms" in state.relit_map_cache - assert "beta terms" in state.relit_map_cache - - -def test_seed_episode_fallback_uses_user_query(monkeypatch: Any) -> None: - seen: list[str] = [] - - def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): - seen.extend(queries) - return {queries[0]: ({"n": 1.0}, {"u": 1.0}, ["u"])} - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", - fake_relight, - ) - state = NavState(doc_id="doc1", query="心血管", task_type="unknown") - plan = RetrievalPlan( - subgoals=[Subgoal(id="s1", need="should not score", retrieval_query="")] - ) - n = seed_episode_map_scores_from_plan( - object(), state, NavConfig(collect_top_k=5), plan - ) - assert n == 1 - assert seen == ["心血管"] - assert state.map_scores == {"n": 1.0} - - -def test_seed_skips_unbound_slot_queries(monkeypatch: Any) -> None: - seen: list[str] = [] - - def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): - seen.extend(queries) - return {q: ({f"n:{q}": 1.0}, {f"u:{q}": 1.0}, [f"u:{q}"]) for q in queries} - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", - fake_relight, - ) - state = NavState(doc_id="", query="user original", task_type="unknown") - plan = RetrievalPlan( - subgoals=[ - Subgoal(id="s1", need="a", retrieval_query="alpha"), - Subgoal( - id="s2", - need="b", - retrieval_query="beta {{s1.entity}}", - depends_on=["s1"], - ), - ] - ) - n = seed_episode_map_scores_from_plan(object(), state, NavConfig(), plan) - assert n == 1 - assert seen == ["alpha"] - assert "alpha" in state.relit_map_cache - assert not any("beta" in q for q in state.relit_map_cache) - - -def test_wave_uses_seeded_cache_and_scores_missing_once(monkeypatch: Any) -> None: - score_calls: list[list[str]] = [] - - def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): - score_calls.append(list(queries)) - return { - q: ({f"n:{q}": 1.0}, {f"u:{q}": 1.0}, [f"u:{q}"]) - for q in queries - } - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", - fake_relight, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._apply_plan_control", - lambda *args, **kwargs: {"done": True}, - ) - - harvest_queries: list[str] = [] - - def fake_harvest(ts: Any, state: NavState, config: NavConfig, **kwargs: Any): - harvest_queries.append(str(kwargs.get("query") or "")) - return type( - "HR", - (), - { - "n_policy_calls": 0, - "visited_section_ids": [], - "max_depth_hit": False, - "reason": "test", - }, - )() - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate._harvest_after_node_filter", - fake_harvest, - ) - - state = NavState(doc_id="", query="user original", task_type="unknown") - plan = RetrievalPlan( - subgoals=[ - Subgoal(id="s1", need="a", retrieval_query="alpha"), - Subgoal(id="s2", need="b", retrieval_query="beta"), - ] - ) - state.retrieval_plan = plan - seed_episode_map_scores_from_plan(object(), state, NavConfig(), plan) - assert score_calls == [["alpha", "beta"]] - - from shared.services.retrieval.nav.nav_orchestrate import execute_plan - - execute_plan(object(), state, NavConfig(), episode_query="user original") - # Wave relight must hit the seeded cache; no second full-corpus scoring. - assert score_calls == [["alpha", "beta"]] - assert harvest_queries == ["alpha", "beta"] - - -def test_run_nav_episode_scores_after_plan(monkeypatch: Any) -> None: - order: list[str] = [] - - class _TS: - def sections_for_doc(self, _doc_id: str) -> list[str]: - return ["sec_root"] - - monkeypatch.setattr( - "shared.services.retrieval.nav._compat.load_llm_env", - lambda: None, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav._compat.require_llm_env", - lambda *_args, **_kwargs: None, - ) - - def fake_plan_query(state: NavState, config: NavConfig, **_kwargs: Any): - order.append("plan") - assert state.map_scores == {} - assert state.unit_scores == {} - return fallback_plan(state.query, reason="test_plan") - - def fake_seed(ts: Any, state: NavState, config: NavConfig, plan: RetrievalPlan): - order.append("score") - assert plan.subgoals - state.map_scores = {"sec_root": 1.0} - state.unit_scores = {"u1": 1.0} - state.highlight_ids = ["u1"] - return 1 - - def fake_execute(ts: Any, state: NavState, config: NavConfig, **_kwargs: Any): - order.append("orch") - assert state.map_scores == {"sec_root": 1.0} - return {"waves": []} - - class _Fill: - scored_chunks = [] - kept_chunks = [] - evidence_text = "" - evidence_chars_actual = 0 - n_chunks_kept = 0 - truncated_last = False - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_plan.plan_query", - fake_plan_query, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate.seed_episode_map_scores_from_plan", - fake_seed, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_orchestrate.execute_plan", - fake_execute, - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_agent.pack_nav_evidence", - lambda *_args, **_kwargs: _Fill(), - ) - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_agent.uses_document_nodes", - lambda _ts: False, - ) - - from shared.services.retrieval.nav.nav_agent import _run_nav_episode_body - - result = _run_nav_episode_body( - None, - "心血管", - doc_id="doc1", - budget_chars=2000, - compose_answer=False, - policy="llm", - config=NavConfig(policy="llm"), - toolspace=_TS(), - ) - assert order == ["plan", "score", "orch"] - assert result.section_ids == ["sec_root"] diff --git a/deprecated/mapnav/tests/shared/test_nav_projection_prod.py b/deprecated/mapnav/tests/shared/test_nav_projection_prod.py deleted file mode 100644 index d03a30bac..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_projection_prod.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Knowhere production projection / config authority tests.""" - -from __future__ import annotations - -import os -from typing import Any - -os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") -os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") -os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") -os.environ.setdefault("S3_ACCESS_KEY_ID", "test") -os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") -os.environ.setdefault("S3_TEMP_PATH", "/tmp") - -from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow -from shared.services.retrieval.nav.nav_projection import ( - _section_summary_for_map, - build_map, -) -from shared.services.retrieval.nav_config import build_nav_config -from shared.services.retrieval.nav_snapshot import build_nav_snapshot - - -def test_build_nav_config_authoritative_for_production() -> None: - cfg = build_nav_config() - assert cfg.mode == "checklist" - - -def test_section_summary_falls_back_to_provider_structure() -> None: - class _FakeTs: - def get_structure(self, section_id: str) -> dict[str, Any]: - assert section_id == "sec_host" - return {"summary": "provider hosted summary", "preview": "preview"} - - text = _section_summary_for_map(_FakeTs(), "sec_host", doc_id="doc_a") - assert text == "provider hosted summary" - - -def test_build_map_inline_summary_uses_provider_when_store_missing() -> None: - root = SectionRow( - section_id="sec_root", - parent_section_id=None, - section_path="Root", - section_title="Root", - section_level=0, - summary="root summary", - sort_order=0, - ) - host = SectionRow( - section_id="sec_host", - parent_section_id="sec_root", - section_path="Chapter 1", - section_title="Chapter 1", - section_level=1, - summary="chapter one summary from db", - sort_order=1, - ) - unit = UnitRow( - chunk_id="chk_text", - section_id="sec_host", - chunk_type="text", - content="body", - sort_order=0, - ) - snap = build_nav_snapshot( - document_titles={"doc_a": "Doc A"}, - sections_by_doc={"doc_a": [root, host]}, - units_by_doc={"doc_a": [unit]}, - chunk_ref_index={ - "chk_text": { - "document_id": "doc_a", - "section_path": "Chapter 1", - "chunk_type": "text", - "file_path": None, - "job_id": "job_1", - } - }, - ) - ts = ProviderToolSpace(snap.provider) - cfg = build_nav_config() - # Force scoped inline-summary path (small actionable map under the limit). - cfg.scope_inline_summary_char_limit = 50_000 - projection = build_map( - ts, - doc_id="doc_a", - query="chapter", - scope="sec_host", - config=cfg, - map_scores={"sec_host": 1.0}, - ) - assert "summary: chapter one summary from db" in projection.text - host_views = [v for v in projection.tree_sections if v.section_id == "sec_host"] - assert host_views - assert "chapter one summary from db" in host_views[0].summary diff --git a/deprecated/mapnav/tests/shared/test_nav_scope_filter.py b/deprecated/mapnav/tests/shared/test_nav_scope_filter.py deleted file mode 100644 index ecf02fc96..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_scope_filter.py +++ /dev/null @@ -1,273 +0,0 @@ -"""Bounded WHERE scope-filter loop tests (mocked policy LLM).""" - -from __future__ import annotations - -import json -from typing import Any, List - -from shared.services.retrieval.nav.nav_knowhere import NamespaceKnowhereProvider -from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace -from shared.services.retrieval.scoring.knowhere_provider import ( - KnowhereProvider, - SectionRow, -) -from shared.services.retrieval.nav.nav_node_filter import node_filter -from shared.services.retrieval.scoring.node_filter_predicates import field_predicate -from shared.services.retrieval.nav.nav_scope_filter import run_scope_filter -from shared.services.retrieval.nav.nav_types import NavConfig - - -def _section( - section_id: str, - parent: str | None, - path: str, - title: str, - *, - level: int, - summary: str = "", - order: int = 0, -) -> SectionRow: - return SectionRow( - section_id=section_id, - parent_section_id=parent, - section_path=path, - section_title=title, - section_level=level, - summary=summary, - sort_order=order, - ) - - -def _ts() -> ProviderToolSpace: - apple = KnowhereProvider( - doc_id="doc_apple", - sections=[ - _section("sec_root_a", None, "Root", "Root", level=0, order=0), - _section( - "sec_q3", - "sec_root_a", - "Q3 Results", - "Q3 Results", - level=1, - summary="Apple quarterly profit", - order=1, - ), - ], - units=(), - ) - filler = KnowhereProvider( - doc_id="doc_other", - sections=[ - _section("sec_root_o", None, "Root", "Root", level=0, order=0), - _section( - "sec_misc", - "sec_root_o", - "Notes", - "Notes", - level=1, - summary="unrelated notes", - order=1, - ), - ], - units=(), - ) - return ProviderToolSpace( - NamespaceKnowhereProvider( - [apple, filler], - titles={"doc_apple": "AAPL 10-K.pdf", "doc_other": "Misc.docx"}, - ) - ) - - -def _cfg(**kwargs: Any) -> NavConfig: - data = { - "enable_node_filter": True, - "filter_max_rounds": 3, - "llm_model": "test-model", - "llm_max_tokens": 256, - } - data.update(kwargs) - return NavConfig.from_dict(data) - - -def _install_script(monkeypatch: Any, replies: List[dict[str, Any]]) -> None: - queue = list(replies) - - def fake_nav_chat(**kwargs: Any) -> dict[str, Any]: - del kwargs - obj = queue.pop(0) if queue else {"action": "fallback", "reason": "empty"} - return {"content": json.dumps(obj)} - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_llm.nav_chat", - fake_nav_chat, - ) - - -def test_widen_relooks_full_map(monkeypatch: Any) -> None: - seen_users: List[str] = [] - queue: List[dict[str, Any]] = [ - {"action": "widen", "reason": "sub-map too narrow"}, - {"action": "done", "decision": "collect_all", "reason": "ok"}, - ] - - def fake_nav_chat(**kwargs: Any) -> dict[str, Any]: - messages = kwargs.get("messages") or [] - seen_users.append(str(messages[-1].get("content") or "")) - obj = queue.pop(0) if queue else {"action": "fallback", "reason": "empty"} - return {"content": json.dumps(obj)} - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_llm.nav_chat", - fake_nav_chat, - ) - out = run_scope_filter( - _ts(), - _cfg(), - query="apple q3 profit", - doc_ids=["doc_apple", "doc_other"], - seed_filter=node_filter([field_predicate("path", ["Q3"])]), - ) - assert out.decision == "collect_all" - assert len(seen_users) == 2 - assert "Full map" not in seen_users[0] - assert "Sub-map" in seen_users[0] - assert "Full map" in seen_users[1] - assert "Last filter" in seen_users[1] - - -def test_zero_hits_widen_then_done(monkeypatch: Any) -> None: - _install_script( - monkeypatch, - [ - { - "action": "filter", - "predicates": [ - {"field": "path", "terms": ["AAPL"], "match": "substring"} - ], - }, - {"action": "done", "decision": "collect_all", "reason": "ok"}, - ], - ) - out = run_scope_filter( - _ts(), - _cfg(), - query="apple q3 profit", - doc_ids=["doc_apple", "doc_other"], - seed_filter=node_filter([field_predicate("path", ["zzz-missing"])]), - ) - assert out.decision == "collect_all" - assert "sec_q3" in out.settled_section_ids - assert out.rounds == 2 - - -def test_zero_hits_policy_fallback(monkeypatch: Any) -> None: - _install_script( - monkeypatch, - [{"action": "fallback", "reason": "cannot widen"}], - ) - out = run_scope_filter( - _ts(), - _cfg(), - query="missing topic", - doc_ids=["doc_apple", "doc_other"], - seed_filter=node_filter([field_predicate("path", ["zzz-missing"])]), - ) - assert out.decision == "fallback" - assert out.reason == "cannot widen" - assert out.rounds == 1 - - -def test_too_many_hits_tighten(monkeypatch: Any) -> None: - _install_script( - monkeypatch, - [ - { - "action": "filter", - "predicates": [ - {"field": "summary", "terms": ["profit"], "match": "substring"} - ], - }, - {"action": "done", "decision": "collect_all", "reason": "tight"}, - ], - ) - out = run_scope_filter( - _ts(), - _cfg(), - query="everything", - doc_ids=["doc_apple", "doc_other"], - seed_filter=node_filter([field_predicate("path", ["Root"])]), - ) - assert out.decision == "collect_all" - assert out.settled_section_ids == ["sec_q3"] - assert out.rounds == 2 - - -def test_max_rounds_hard_stop(monkeypatch: Any) -> None: - _install_script( - monkeypatch, - [ - { - "action": "filter", - "predicates": [ - {"field": "path", "terms": ["zzz"], "match": "substring"} - ], - }, - { - "action": "filter", - "predicates": [ - {"field": "path", "terms": ["still-missing"], "match": "substring"} - ], - }, - ], - ) - out = run_scope_filter( - _ts(), - _cfg(filter_max_rounds=2), - query="no hits", - doc_ids=["doc_apple"], - seed_filter=node_filter([field_predicate("path", ["absent"])]), - ) - assert out.decision == "fallback" - assert out.rounds == 2 - assert out.reason == "max_rounds" - - -def test_cardinality_drives_scoped_harvest(monkeypatch: Any) -> None: - _install_script( - monkeypatch, - [{"action": "done", "reason": "keep"}], - ) - out = run_scope_filter( - _ts(), - _cfg(), - query="apple", - doc_ids=["doc_apple", "doc_other"], - seed_filter=node_filter([field_predicate("path", ["AAPL"])]), - ) - assert out.decision == "scoped_harvest" - assert out.rounds == 1 - assert "sec_q3" in out.settled_section_ids - - -def test_disabled_skips_policy(monkeypatch: Any) -> None: - called = {"n": 0} - - def boom(**kwargs: Any) -> dict[str, Any]: - called["n"] += 1 - raise AssertionError("policy should not run") - - monkeypatch.setattr( - "shared.services.retrieval.nav.nav_llm.nav_chat", - boom, - ) - out = run_scope_filter( - _ts(), - _cfg(enable_node_filter=False), - query="apple", - doc_ids=["doc_apple"], - seed_filter=node_filter([field_predicate("path", ["AAPL"])]), - ) - assert out.decision == "fallback" - assert out.reason == "disabled" - assert called["n"] == 0 diff --git a/deprecated/mapnav/tests/shared/test_nav_snapshot.py b/deprecated/mapnav/tests/shared/test_nav_snapshot.py deleted file mode 100644 index 76911ac19..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_snapshot.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Unit tests for map-nav snapshot assembly (no DB).""" - -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") - -import pytest - -from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow -from shared.services.retrieval.nav_snapshot import build_nav_snapshot - - -def test_build_nav_snapshot_keeps_original_section_path_after_root_remount() -> None: - """chunk_ref_index must keep DB Root path even if provider remounts assets.""" - root = SectionRow( - section_id="sec_root", - parent_section_id=None, - section_path="Root", - section_title="Root", - section_level=0, - summary="", - sort_order=0, - ) - host = SectionRow( - section_id="sec_host", - parent_section_id="sec_root", - section_path="Chapter 1", - section_title="Chapter 1", - section_level=1, - summary="host", - sort_order=1, - ) - text = UnitRow( - chunk_id="chk_text", - section_id="sec_host", - chunk_type="text", - content="body", - sort_order=0, - metadata={"connect_to": [{"target": "chk_img"}]}, - ) - image = UnitRow( - chunk_id="chk_img", - section_id="sec_root", - chunk_type="image", - content="img/a.png", - sort_order=1, - file_path="img/a.png", - ) - ref_index = { - "chk_text": { - "document_id": "doc_a", - "section_path": "Chapter 1", - "chunk_type": "text", - "file_path": None, - "job_id": "job_1", - }, - "chk_img": { - "document_id": "doc_a", - "section_path": "Root", # DB original - "chunk_type": "image", - "file_path": "img/a.png", - "job_id": "job_1", - }, - } - snap = build_nav_snapshot( - document_titles={"doc_a": "Doc A"}, - sections_by_doc={"doc_a": [root, host]}, - units_by_doc={"doc_a": [text, image]}, - chunk_ref_index=ref_index, - ) - - assert snap.document_ids == ["doc_a"] - assert snap.chunk_ref_index["chk_img"]["section_path"] == "Root" - # Provider remounted the image onto the host section for navigation. - host_units = list(snap.provider.self_units("sec_host")) - assert any(u.chunk_id == "chk_img" for u in host_units) - root_units = list(snap.provider.self_units("sec_root")) - assert not any(u.chunk_id == "chk_img" for u in root_units) - - -def test_build_nav_snapshot_rejects_empty_corpus() -> None: - with pytest.raises(ValueError, match="at least one"): - build_nav_snapshot( - document_titles={}, - sections_by_doc={}, - units_by_doc={}, - chunk_ref_index={}, - ) diff --git a/deprecated/mapnav/tests/shared/test_nav_stamp.py b/deprecated/mapnav/tests/shared/test_nav_stamp.py deleted file mode 100644 index 539421cb1..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_stamp.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Unit test for AgentStep stamp_step_detail.""" - -from __future__ import annotations - -import time - -from shared.services.retrieval.nav.nav_token_budget import ( - nav_token_episode, - record_episode_tokens, - stamp_step_detail, -) - - -def test_stamp_step_detail_tracks_delta_and_elapsed() -> None: - with nav_token_episode(): - t0 = time.perf_counter() - first = stamp_step_detail({"action": "a"}, t0=t0) - assert first["tokens_used_total"] == 0 - assert first["tokens_used_delta"] == 0 - assert first["elapsed_ms"] >= 0 - record_episode_tokens({"total_tokens": 40}) - second = stamp_step_detail({"action": "b"}) - assert second["tokens_used_total"] == 40 - assert second["tokens_used_delta"] == 40 - record_episode_tokens({"total_tokens": 10}) - third = stamp_step_detail({"action": "c"}) - assert third["tokens_used_total"] == 50 - assert third["tokens_used_delta"] == 10 diff --git a/deprecated/mapnav/tests/shared/test_nav_trace_map.py b/deprecated/mapnav/tests/shared/test_nav_trace_map.py deleted file mode 100644 index 55a2e0500..000000000 --- a/deprecated/mapnav/tests/shared/test_nav_trace_map.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Unit tests for map-nav AgentStep → decision_trace mapping.""" - -from __future__ import annotations - -import os -from types import SimpleNamespace - -os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") -os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") -os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") -os.environ.setdefault("S3_ACCESS_KEY_ID", "test") -os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") -os.environ.setdefault("S3_TEMP_PATH", "/tmp") - -from shared.services.retrieval.trace.mapnav import ( - build_decision_trace, - episode_selected_doc_ids, - episode_selected_paths, - episode_token_count, - episode_workflow_plan, -) - - -def _step(action: str, detail: dict, step_idx: int = 1) -> SimpleNamespace: - return SimpleNamespace(step_idx=step_idx, action=action, detail=detail) - - -def test_build_decision_trace_maps_three_layers_and_terminal() -> None: - episode = SimpleNamespace( - stop_reason="completed", - evidence_chars_actual=1200, - steps=[ - _step( - "query_plan", - { - "reason": "plan it", - "n_subgoals": 1, - "fallback": False, - "plan": { - "subgoals": [{"id": "s1"}], - "coverage_checklist": [{"id": "c1", "fact": "x"}], - "map_coverage": "sufficient", - "raw": '{"ok":true}', - "reason": "plan it", - }, - "token_limit": 100000, - "tokens_used_total": 100, - "tokens_used_delta": 100, - "elapsed_ms": 10, - }, - ), - _step( - "harvest", - { - "subgoal_id": "s1", - "scope": "sec_a", - "depth": 0, - "collect_ids": ["C1"], - "dispatch_ids": [], - "search_assets": [], - "confidence": {"C1": 0.9}, - "reason": "hit", - "projection_chars": 200, - "legal_actions_preview": ["C1 collect"], - "visible_section_ids": ["sec_a"], - "model": "deepseek-v4-flash", - "raw": "{}", - "token_limit": 100000, - "tokens_used_total": 250, - "tokens_used_delta": 150, - "elapsed_ms": 20, - }, - step_idx=2, - ), - _step( - "harvest", - { - "subgoal_id": "s1", - "scope": "sec_b", - "depth": 1, - "collect_ids": ["C2"], - "dispatch_ids": [], - "search_assets": [], - "reason": "child", - "token_limit": 100000, - "tokens_used_total": 300, - "tokens_used_delta": 50, - "elapsed_ms": 5, - }, - step_idx=3, - ), - _step( - "plan_control", - { - "global": "done", - "reason": "covered", - "subgoals": {"s1": {"decision": "accept", "note": "ok"}}, - "token_limit": 100000, - "tokens_used_total": 400, - "tokens_used_delta": 100, - "elapsed_ms": 8, - }, - step_idx=4, - ), - ], - ) - steps = build_decision_trace(episode, evidence_char_budget=12000, n_refs=2) - phases = [s.phase for s in steps] - assert phases[0] == "plan" - assert phases[1] == "harvest" - assert phases[2] == "harvest" - assert phases[3] == "plan_control" - assert phases[-1] == "terminal" - assert steps[0].agent == "planner" - assert steps[0].decision["action"] == "plan_query" - assert steps[0].result["coverage_checklist"][0]["id"] == "c1" - assert steps[0].budget["tokens_used_delta"] == 100 - assert steps[1].observation["projection_chars"] == 200 - assert steps[1].decision["collect_ids"] == ["C1"] - assert steps[2].parent_step_index == steps[1].step_index - assert steps[3].decision["action"] == "done" - assert steps[-1].result["n_refs"] == 2 - assert steps[-1].result["layer_llm_steps"]["planner"] >= 1 - assert steps[-1].result["layer_llm_steps"]["control"] >= 1 - - -def test_node_filter_steps_map_and_count_tokens() -> None: - episode = SimpleNamespace( - stop_reason="completed", - evidence_chars_actual=0, - steps=[ - _step( - "node_filter", - { - "round": 1, - "predicates": [ - {"field": "path", "terms": ["AAPL"], "match": "substring"} - ], - "fields": ["path"], - "cardinality": 2, - "action": "done", - "decision": "collect_all", - "reason": "hits", - "matched_section_ids": ["sec_q3"], - "token_limit": 100000, - "tokens_used_total": 80, - "tokens_used_delta": 80, - "elapsed_ms": 12, - }, - ) - ], - ) - steps = build_decision_trace(episode, evidence_char_budget=12000, n_refs=0) - assert steps[0].phase == "node_filter" - assert steps[0].observation["cardinality"] == 2 - assert steps[0].observation["fields"] == ["path"] - assert steps[0].decision["action"] == "done" - assert steps[0].result["decision"] == "collect_all" - assert steps[0].budget["tokens_used_delta"] == 80 - assert steps[0].budget["token_limit"] == 100000 - assert steps[-1].result["layer_llm_steps"]["harvest"] >= 1 - - -def test_episode_helpers() -> None: - episode = SimpleNamespace( - steps=[ - _step( - "query_plan", - { - "plan": {"subgoals": [{"id": "s1"}], "raw": "x"}, - "tokens_used_total": 12, - }, - ), - _step("harvest", {"tokens_used_total": 40}), - ] - ) - refs = [ - {"document_id": "d1", "section_path": "A / B", "chunk_id": "c1"}, - {"document_id": "d1", "section_path": "A / B", "chunk_id": "c2"}, - {"document_id": "d2", "section_path": "C", "chunk_id": "c3"}, - ] - assert episode_workflow_plan(episode)["subgoals"][0]["id"] == "s1" - assert episode_token_count(episode) == 40 - assert episode_selected_paths(episode, refs) == ["A / B", "C"] - assert episode_selected_doc_ids(refs) == ["d1", "d2"] diff --git a/deprecated/mapnav/trace_mapnav.py b/deprecated/mapnav/trace_mapnav.py deleted file mode 100644 index 3d97659d7..000000000 --- a/deprecated/mapnav/trace_mapnav.py +++ /dev/null @@ -1,475 +0,0 @@ -"""Map map-nav ``AgentStep`` list → public ``decision_trace``. - -Replacement for the former stub: one DecisionTraceStep per stamped AgentStep -(plus a terminal row), shaped for PLANNER / HARVEST / CONTROL. -""" - -from __future__ import annotations - -from typing import Any, Optional, Sequence - -from shared.services.retrieval.nav_config import ( - MAPNAV_EVIDENCE_CHARS, - MAPNAV_MODEL, - MAPNAV_TRACE_RAW_CHARS, -) -from shared.services.retrieval.trace.types import DecisionTraceStep - - -def _detail(step: Any) -> dict[str, Any]: - raw = getattr(step, "detail", None) - return dict(raw) if isinstance(raw, dict) else {} - - -def _clip_raw(text: Any) -> str: - s = str(text or "") - limit = int(MAPNAV_TRACE_RAW_CHARS) - if limit > 0 and len(s) > limit: - return s[:limit] - return s - - -def _budget_from_detail( - detail: dict[str, Any], - *, - evidence_chars_used: int, - evidence_char_budget: int, -) -> dict[str, Any]: - token_limit = int(detail.get("token_limit") or 0) - used_total = int(detail.get("tokens_used_total") or 0) - used_delta = int(detail.get("tokens_used_delta") or 0) - remaining = max(0, token_limit - used_total) if token_limit > 0 else None - out: dict[str, Any] = { - "token_limit": token_limit, - "tokens_used_total": used_total, - "tokens_used_delta": used_delta, - "evidence_chars_used": int(evidence_chars_used), - "evidence_char_budget": int(evidence_char_budget), - } - if remaining is not None: - out["remaining"] = remaining - return out - - -def _elapsed_ms(detail: dict[str, Any]) -> Optional[int]: - if "elapsed_ms" in detail and detail.get("elapsed_ms") is not None: - try: - return int(detail.get("elapsed_ms") or 0) - except (TypeError, ValueError): - return None - if "seconds" in detail and detail.get("seconds") is not None: - try: - return int(round(float(detail.get("seconds") or 0.0) * 1000.0)) - except (TypeError, ValueError): - return None - return None - - -def _map_one( - *, - action: str, - detail: dict[str, Any], - step_index: int, - parent_step_index: Optional[int], - evidence_chars_used: int, - evidence_char_budget: int, -) -> Optional[DecisionTraceStep]: - budget = _budget_from_detail( - detail, - evidence_chars_used=evidence_chars_used, - evidence_char_budget=evidence_char_budget, - ) - elapsed = _elapsed_ms(detail) - scope = detail.get("scope") - if scope is not None: - scope = str(scope) - - if action == "query_plan": - raw_plan = detail.get("plan") - plan_payload = raw_plan if isinstance(raw_plan, dict) else {} - return DecisionTraceStep( - step_index=step_index, - agent="planner", - phase="plan", - parent_step_index=parent_step_index, - scope=scope, - observation={ - "llm_raw": _clip_raw( - plan_payload.get("raw") or detail.get("llm_raw") or detail.get("raw") - ), - }, - decision={ - "action": "plan_query", - "reason": detail.get("reason") or plan_payload.get("reason") or "", - }, - result={ - "status": "fallback" if detail.get("fallback") else "ok", - "subgoals": plan_payload.get("subgoals") or [], - "coverage_checklist": plan_payload.get("coverage_checklist") or [], - "map_coverage": plan_payload.get("map_coverage"), - "n_subgoals": detail.get("n_subgoals"), - "reason": detail.get("reason") or plan_payload.get("reason") or "", - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "plan_orchestrate": - return DecisionTraceStep( - step_index=step_index, - agent="planner", - phase="plan", - parent_step_index=parent_step_index, - scope=scope, - observation={}, - decision={"action": "orchestrate"}, - result={ - "status": "ok", - "n_waves": detail.get("n_waves") - or len(detail.get("waves") or []), - "satisfied": detail.get("satisfied"), - "attempted": detail.get("attempted"), - "dropped": detail.get("dropped"), - "waves": detail.get("waves"), - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "plan_wave": - return DecisionTraceStep( - step_index=step_index, - agent="orchestrator", - phase="plan_wave", - parent_step_index=parent_step_index, - scope=scope, - observation={}, - decision={ - "action": "plan_wave", - "wave": detail.get("wave"), - "subgoal_ids": detail.get("subgoal_ids") - or detail.get("ready") - or [], - }, - result={ - "status": "ok", - "subgoals": detail.get("subgoals") or detail.get("results"), - "chars_used": detail.get("chars_used"), - "plan_control": detail.get("plan_control"), - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "harvest": - return DecisionTraceStep( - step_index=step_index, - agent="navigator", - phase="harvest", - parent_step_index=parent_step_index, - scope=scope, - observation={ - "visible_section_ids": detail.get("visible_section_ids") or [], - "legal_actions_preview": detail.get("legal_actions_preview") or [], - "projection_chars": detail.get("projection_chars"), - "scope": scope, - "depth": detail.get("depth"), - "llm_raw": _clip_raw(detail.get("raw") or detail.get("llm_raw")), - }, - decision={ - "action": "harvest", - "collect_ids": detail.get("collect_ids") or [], - "dispatch_ids": detail.get("dispatch_ids") or [], - "search_assets": detail.get("search_assets") or [], - "confidence": detail.get("confidence"), - "reason": detail.get("reason") or "", - }, - result={ - "status": "ok", - "subgoal_id": detail.get("subgoal_id"), - "collect_section_ids": detail.get("collect_section_ids") or [], - "n_added": detail.get("n_added"), - "fallback_used": detail.get("fallback_used"), - "model": detail.get("model") or MAPNAV_MODEL, - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "node_filter": - step_action = str(detail.get("action") or detail.get("decision") or "filter") - return DecisionTraceStep( - step_index=step_index, - agent="navigator", - phase="node_filter", - parent_step_index=parent_step_index, - scope=scope, - observation={ - "predicates": detail.get("predicates") or [], - "fields": detail.get("fields") or [], - "cardinality": detail.get("cardinality"), - "failed_predicates": detail.get("failed_predicates") or [], - "matched_section_ids": detail.get("matched_section_ids") or [], - "round": detail.get("round"), - }, - decision={ - "action": step_action, - "reason": detail.get("reason") or "", - }, - result={ - "status": "fallback" if step_action == "fallback" else "ok", - "cardinality": detail.get("cardinality"), - "decision": detail.get("decision") or "", - "reason": detail.get("reason") or "", - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "search_assets": - return DecisionTraceStep( - step_index=step_index, - agent="asset_inspector", - phase="asset_search", - parent_step_index=parent_step_index, - scope=scope, - observation={ - "candidates": detail.get("candidates") - or detail.get("requests") - or [], - }, - decision={ - "action": "search_assets", - "asset_ids": detail.get("asset_ids") or [], - }, - result={ - "status": "ok", - "n_added": detail.get("n_added"), - "subgoal_id": detail.get("subgoal_id"), - "requests": detail.get("requests"), - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "plan_control": - return DecisionTraceStep( - step_index=step_index, - agent="controller", - phase="plan_control", - parent_step_index=parent_step_index, - scope=scope, - observation={ - "coverage": detail.get("coverage") - or detail.get("checklist") - or detail.get("subgoals"), - }, - decision={ - "action": detail.get("global") or "continue", - "subgoals": detail.get("subgoals") or {}, - "reason": detail.get("reason") or "", - }, - result={ - "status": "ok", - "global": detail.get("global"), - "reason": detail.get("reason") or "", - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "replan": - raw_plan = detail.get("plan") - plan_payload = raw_plan if isinstance(raw_plan, dict) else {} - return DecisionTraceStep( - step_index=step_index, - agent="controller", - phase="plan_control", - parent_step_index=parent_step_index, - scope=scope, - observation={}, - decision={ - "action": "replan", - "reason": detail.get("reason") or "", - }, - result={ - "status": "ok", - "plan": plan_payload or detail.get("plan"), - "n_subgoals": detail.get("n_subgoals"), - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "slot_extract": - return DecisionTraceStep( - step_index=step_index, - agent="controller", - phase="plan_control", - parent_step_index=parent_step_index, - scope=scope, - observation={ - "slots_requested": detail.get("slots_requested") or [], - "llm_raw": _clip_raw(detail.get("raw") or detail.get("llm_raw")), - }, - decision={"action": "slot_extract"}, - result={ - "status": "ok", - "subgoal_id": detail.get("subgoal_id"), - "slots": detail.get("slots") or detail.get("extracted") or {}, - }, - budget=budget, - elapsed_ms=elapsed, - ) - - if action == "compose_answer": - return DecisionTraceStep( - step_index=step_index, - agent="composer", - phase="compose", - parent_step_index=parent_step_index, - scope=scope, - observation={}, - decision={"action": "compose_answer"}, - result={"status": "ok", **{ - k: detail.get(k) - for k in ("chars", "model", "truncated") - if k in detail - }}, - budget=budget, - elapsed_ms=elapsed, - ) - - return None - - -def build_decision_trace( - episode: Any, - *, - evidence_char_budget: int = MAPNAV_EVIDENCE_CHARS, - n_refs: int = 0, -) -> list[DecisionTraceStep]: - """Map episode AgentSteps → DecisionTraceStep list (incl. terminal).""" - evidence_chars_used = int(getattr(episode, "evidence_chars_actual", 0) or 0) - steps_in: Sequence[Any] = list(getattr(episode, "steps", None) or ()) - out: list[DecisionTraceStep] = [] - harvest_parent_by_depth: dict[int, int] = {} - layer_counts = {"planner": 0, "harvest": 0, "control": 0} - - for raw_step in steps_in: - action = str(getattr(raw_step, "action", "") or "").strip() - detail = _detail(raw_step) - parent: Optional[int] = None - if action == "harvest": - try: - depth = int(detail.get("depth") or 0) - except (TypeError, ValueError): - depth = 0 - if depth > 0: - parent = harvest_parent_by_depth.get(depth - 1) - mapped = _map_one( - action=action, - detail=detail, - step_index=len(out), - parent_step_index=parent, - evidence_chars_used=evidence_chars_used, - evidence_char_budget=evidence_char_budget, - ) - if mapped is None: - continue - out.append(mapped) - if action == "harvest": - try: - depth = int(detail.get("depth") or 0) - except (TypeError, ValueError): - depth = 0 - harvest_parent_by_depth[depth] = mapped.step_index - if mapped.phase == "plan": - layer_counts["planner"] += 1 - elif mapped.phase in {"harvest", "plan_wave", "asset_search", "node_filter"}: - layer_counts["harvest"] += 1 - elif mapped.phase == "plan_control": - layer_counts["control"] += 1 - - stop = str(getattr(episode, "stop_reason", "") or "completed") - last_budget = out[-1].budget if out and out[-1].budget else { - "token_limit": 0, - "tokens_used_total": 0, - "tokens_used_delta": 0, - "evidence_chars_used": evidence_chars_used, - "evidence_char_budget": int(evidence_char_budget), - } - out.append( - DecisionTraceStep( - step_index=len(out), - agent="retrieval_agent", - phase="terminal", - scope="mapnav", - observation={ - "router_used": "mapnav", - "n_nav_steps": len(steps_in), - "evidence_chars": evidence_chars_used, - }, - decision={"action": "complete", "args": {}, "reason": stop}, - result={ - "status": "ok", - "stop_reason": stop, - "evidence_chars": evidence_chars_used, - "n_refs": int(n_refs), - "layer_llm_steps": dict(layer_counts), - }, - budget=dict(last_budget), - elapsed_ms=None, - ) - ) - return out - - -def episode_token_count(episode: Any) -> int: - """Max stamped tokens_used_total on AgentSteps (episode counter is gone after return).""" - best = 0 - for step in getattr(episode, "steps", None) or (): - detail = _detail(step) - try: - best = max(best, int(detail.get("tokens_used_total") or 0)) - except (TypeError, ValueError): - continue - return best - - -def episode_workflow_plan(episode: Any) -> Optional[dict[str, Any]]: - """Nav plan dict for ``retrieval_runs.workflow_plan``.""" - for step in getattr(episode, "steps", None) or (): - if str(getattr(step, "action", "") or "") != "query_plan": - continue - detail = _detail(step) - plan = detail.get("plan") - if isinstance(plan, dict): - return plan - return None - - -def episode_selected_paths( - episode: Any, - refs: Sequence[dict[str, Any]], -) -> list[str]: - """Section paths from resolved refs (DB-original paths on refs).""" - paths: list[str] = [] - seen: set[str] = set() - for ref in refs or (): - path = str(ref.get("section_path") or "").strip() - if path and path not in seen: - seen.add(path) - paths.append(path) - return paths - - -def episode_selected_doc_ids(refs: Sequence[dict[str, Any]]) -> list[str]: - docs: list[str] = [] - seen: set[str] = set() - for ref in refs or (): - doc = str(ref.get("document_id") or "").strip() - if doc and doc not in seen: - seen.add(doc) - docs.append(doc) - return docs diff --git a/docs/adr/0005-stream-retrieval-progress-over-sse.md b/docs/adr/0005-stream-retrieval-progress-over-sse.md index dfaafbd85..3333314f0 100644 --- a/docs/adr/0005-stream-retrieval-progress-over-sse.md +++ b/docs/adr/0005-stream-retrieval-progress-over-sse.md @@ -7,10 +7,10 @@ Accepted ## Context Online Brain users currently wait for the complete Retrieval response while -map-nav planning, searching, and source review run. Knowhere owns retrieval and -evidence, while answer generation belongs to downstream clients. A streaming -contract must improve perceived latency without exposing chain-of-thought or -making partial citations authoritative. +classic or agent_explore planning, searching, and source review run. Knowhere +owns retrieval and evidence, while answer generation belongs to downstream +clients. A streaming contract must improve perceived latency without exposing +chain-of-thought or making partial citations authoritative. ## Decision @@ -23,18 +23,18 @@ include only safe aggregate counts. The stream terminates with a versioned envelope containing either the existing retrieval response as the authoritative result or a typed, user-safe failure (`failed`, `cancelled`, or `no_results`). -The synchronous map-nav engine remains intact and publishes sanitized progress -through an optional step callback bridged to the SSE route by an async queue. -Answer-token streaming remains downstream. The endpoint sends heartbeats, -disables proxy buffering, and uses per-connection event IDs without promising -replay in the first version. +Classic and agent_explore publish sanitized progress through an optional step +callback bridged to the SSE route by an async queue. Answer-token streaming +remains downstream. The endpoint sends heartbeats, disables proxy buffering, +and uses per-connection event IDs without promising replay in the first +version. Retrieval duration is measured from the execution plan's start through final public-result assembly. The same duration definition is used for persisted `retrieval_runs.latency_ms` and latency aggregates; cache hits are recorded too. Trace persistence must receive the execution start timestamp rather than starting its own timer after navigation and hydration. Timing and terminal -outcomes must cover classic, map-nav, small-corpus, cache-hit, failed, and +outcomes must cover classic, agent_explore, small-corpus, cache-hit, failed, and cancelled routes. `retrieval_runs` is the ledger for all of those routes and stores explicit route, cache, latency, and terminal-status fields. diff --git a/docs/design/retrieval-heart-vessel-trace-optimization-plan.md b/docs/design/retrieval-heart-vessel-trace-optimization-plan.md index f4521260b..0b8ac42cd 100644 --- a/docs/design/retrieval-heart-vessel-trace-optimization-plan.md +++ b/docs/design/retrieval-heart-vessel-trace-optimization-plan.md @@ -1,5 +1,8 @@ # Retrieval optimization plan from the fresh `心血管` trace +Status: historical. The map-nav route named below was later removed. Live +routes are classic map-unit BM25 and agent_explore. + ## Scope and evidence The detailed stage breakdown below is based on a newly issued production diff --git a/docs/design/retrieval-serving-index-rollout-runbook.md b/docs/design/retrieval-serving-index-rollout-runbook.md index 2d91c97b8..2efae7765 100644 --- a/docs/design/retrieval-serving-index-rollout-runbook.md +++ b/docs/design/retrieval-serving-index-rollout-runbook.md @@ -1,5 +1,8 @@ # Retrieval Serving Index Rollout Runbook +Status: historical. The map-nav smoke named below was later removed. Live +routes are classic map-unit BM25 and agent_explore. + ## Scope This runbook rolls out the retrieval serving-index performance changes without diff --git a/docs/design/retrieval-serving-snapshot-cache.md b/docs/design/retrieval-serving-snapshot-cache.md deleted file mode 100644 index 36f23f9df..000000000 --- a/docs/design/retrieval-serving-snapshot-cache.md +++ /dev/null @@ -1,48 +0,0 @@ -# Retrieval serving snapshot cache - -Map-nav reads a namespace routing snapshot that contains section/chunk -relationships, ordering, chunk types, connection references, and remount -ownership. It intentionally does not contain chunk body text; final hydration -still reads pinned revision content from PostgreSQL. - -## Persisted formats - -- Namespace snapshots written by current publication code use format version 2 - and contain routing metadata only. -- The reader remains compatible with version 1 snapshots so existing rows can - be served during rollout. No unconditional namespace rebuild is required - solely because the reader was upgraded. -- A missing, corrupt, stale, or generation-mismatched snapshot falls back to - the exact manifest/table loading path. - -## Redis cache - -The map-nav reader caches the compressed snapshot bytes, not a decoded Python -dictionary. Keys are scoped by user, normalized namespace, and serving -generation: - -```text -retrieval:snapshot:v2:{user_id}:{namespace}:g{generation} -``` - -The cache TTL is one hour. Redis errors, misses, and invalid blobs are -non-fatal: PostgreSQL remains the source of truth and the reader repopulates -Redis after a successful database read. Binary snapshot operations use a Redis -connection with response decoding disabled; normal JSON Redis operations keep -their existing text-decoding connection. - -## Generation coherence - -Retrieval carries one captured revision set and namespace generation through -snapshot loading, map-nav, reference resolution, and final hydration. A -generation lookup failure is treated as inability to establish coherence and -uses the fallback path. A generation mismatch is never served as a valid -snapshot. - -## Diagnostics - -Set `RETRIEVAL_SNAPSHOT_TIMING=1` for detailed snapshot decode timings during -local or staging diagnosis. The route CPU metric is process-level CPU time for -the map-nav route. `ru_maxrss` is a process high-water mark, not a request-level -memory peak, and must not be interpreted as one. - diff --git a/docs/design/retrieval-streaming-sse.md b/docs/design/retrieval-streaming-sse.md index cc8e97468..2ff09b830 100644 --- a/docs/design/retrieval-streaming-sse.md +++ b/docs/design/retrieval-streaming-sse.md @@ -117,10 +117,9 @@ are omitted when not yet known. The execution plan already starts a monotonic timer before cache lookup and logs elapsed time after the route. However, `TraceRecorder` currently starts a -second timer in its constructor, and the map-nav route constructs it only -after navigation, reference resolution, and result assembly. Persisted -`retrieval_runs.latency_ms` and its aggregates therefore under-report retrieval -latency and mostly measure trace flush time. +second timer in its constructor. Persisted `retrieval_runs.latency_ms` and its +aggregates therefore under-report retrieval latency and mostly measure trace +flush time. The canonical `Retrieval Duration` is the execution-plan timer from retrieval start through final public-result assembly. It includes cache lookup and @@ -150,13 +149,11 @@ its latency update must still use the captured execution timestamps and must not extend the user-visible retrieval duration with an unbounded database flush. -Because the map-nav route deliberately rolls back its request session before -the synchronous LLM episode, a trace row must not be created in that session -before the rollback. Capture the execution start immediately, then create or -update the trace record at the terminal persistence point with the captured -start and an explicit end/duration supplied by the execution plan (or use a -separate trace session). `TraceRecorder.complete()` must not silently choose a -later local constructor time or include its own flush duration. +Capture the execution start immediately, then create or update the trace +record at the terminal persistence point with the captured start and an +explicit end/duration supplied by the execution plan (or use a separate +trace session). `TraceRecorder.complete()` must not silently choose a later +local constructor time or include its own flush duration. ## Operational requirements diff --git a/docs/vlm-usage-points.md b/docs/vlm-usage-points.md index 157565ed5..c1b2efe7e 100644 --- a/docs/vlm-usage-points.md +++ b/docs/vlm-usage-points.md @@ -70,7 +70,6 @@ Table / body summaries on this track are text LLM, not VLM. - Scanned-page text: local RapidOCR - TOC source choose, cross-page table continuity, title-tree refine, Excel/DOCX table summary: text LLM -- Archived map-nav image filter: not on the live path ## Debug diff --git a/packages/shared-python/pyproject.toml b/packages/shared-python/pyproject.toml index f94bb9ebf..8baf9c07c 100644 --- a/packages/shared-python/pyproject.toml +++ b/packages/shared-python/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ "psycopg2-binary==2.9.12", "alembic==1.13.1", "sqlalchemy-utils==0.41.1", - "pgvector==0.2.4", + "pgvector>=0.4.1", # Redis "redis==5.3.1", diff --git a/packages/shared-python/shared/core/celery_app.py b/packages/shared-python/shared/core/celery_app.py index d4b5e39b6..5d24d80a4 100644 --- a/packages/shared-python/shared/core/celery_app.py +++ b/packages/shared-python/shared/core/celery_app.py @@ -97,6 +97,13 @@ def get_unique_node_name() -> str: "app.core.tasks.document_ingestion_tasks.*": { "queue": "document_ingestion_medium" }, + "app.core.tasks.enrichment_tasks.build_page_memory_structure": { + "queue": "document_ingestion_medium" + }, + "app.core.tasks.enrichment_tasks.sweep_progressive_enrichments": { + "queue": "default" + }, + "app.core.tasks.enrichment_tasks.*": {"queue": "document_ingestion_low"}, "app.core.tasks.kb_tasks.*": {"queue": "kb_medium"}, # Webhook orphan recovery "app.core.tasks.webhook_tasks.recover_orphaned_webhooks": {"queue": "default"}, @@ -113,6 +120,10 @@ def get_unique_node_name() -> str: "task": "app.core.tasks.stale_job_sweeper.expire_stale_jobs", "schedule": 1800.0, # Every 30 minutes }, + "sweep-progressive-enrichments": { + "task": "app.core.tasks.enrichment_tasks.sweep_progressive_enrichments", + "schedule": 300.0, # Every 5 minutes + }, }, ) diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index 3dc00f5cf..7379ffd41 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -44,6 +44,36 @@ class AIConfig(BaseModel): "Empty follows IMAGE_MODEL." ), ) + VERCEL_KEY: str = Field( + default="", + description="Vercel AI Gateway API key for page-image and query embeddings.", + ) + EMBEDDING_MODEL: str = Field( + default="google/gemini-embedding-2", + description=( + "Vercel AI Gateway embedding model producing shared-space " + "page-image and query vectors for scanned page retrieval." + ), + ) + EMBEDDING_DIM: int = Field( + default=3072, + description="Dimension of page-image and query embeddings.", + ) + EMBEDDING_ENDPOINT: str = Field( + default="https://ai-gateway.vercel.sh/v1/embeddings", + description="Vercel AI Gateway OpenAI-compatible embeddings endpoint.", + ) + VISUAL_RETRIEVAL_ENABLED: bool = Field( + default=True, + description=( + "Enable pure-visual M0 page embeddings for scanned progressive " + "documents and hybrid visual+lexical retrieval." + ), + ) + PAGE_MEMORY_EMBED_CONCURRENCY: int = Field( + default=8, + description="Local per-job page embedding concurrency for scanned M0 documents.", + ) # Runtime LLM controls. LLM_MOCK_ENABLED: bool = Field( diff --git a/packages/shared-python/shared/core/constants/business.py b/packages/shared-python/shared/core/constants/business.py index c17d5d213..f80b1ad51 100644 --- a/packages/shared-python/shared/core/constants/business.py +++ b/packages/shared-python/shared/core/constants/business.py @@ -11,10 +11,6 @@ class BusinessConstants: "N_EPOCHS": 3, "THRESHOLD": 0.5, "OCR_TIMEOUT": 30.0, - "HYBRID_SEARCH": True, - "USE_LOCAL_LLM": False, - "USE_LOCAL_ENCODER": False, - "USE_LOCAL_RERANKER": False, "USE_STOPWORDS": False, "LOCAL_SUMMARY": False, "API_NAME": "ds_api", @@ -33,7 +29,6 @@ class BusinessConstants: # Model-related defaults. MIN_CONFIDENCE_THRESHOLD = 0.05 HIGH_IOU_THRESHOLD = 0.9 - DEFAULT_EMBEDDING_DIM = 1024 DEFAULT_BATCH_SIZE = 32 DEFAULT_EPOCHS = 3 DEFAULT_LEARNING_RATE = 0.001 diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index 7426b8262..986696b2e 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -16,6 +16,8 @@ DocumentMapUnit, DocumentMapUnitIndex, DocumentMapUnitToken, + DocumentPageEmbedding, + DocumentScopeEnrichment, DocumentSection, GraphEdge, GraphNode, @@ -59,9 +61,11 @@ "Document", "DocumentSection", "DocumentChunk", + "DocumentScopeEnrichment", "DocumentMapUnit", "DocumentMapUnitIndex", "DocumentMapUnitToken", + "DocumentPageEmbedding", "DocumentPagePlan", "DemoMaterialization", "GraphNode", diff --git a/packages/shared-python/shared/models/database/demo_corpus.py b/packages/shared-python/shared/models/database/demo_corpus.py index 8e9231083..004acc475 100644 --- a/packages/shared-python/shared/models/database/demo_corpus.py +++ b/packages/shared-python/shared/models/database/demo_corpus.py @@ -6,6 +6,7 @@ from typing import Any, Dict, Optional from uuid import uuid4 +from pgvector.sqlalchemy import HALFVEC from sqlalchemy import ( JSON, CheckConstraint, @@ -270,9 +271,9 @@ class DemoDocumentChunk(Base): class DemoDocumentMapUnit(Base): """Persisted lexical map unit for one document revision. - These rows are a derived index of the exact leaf and interstitial units - used by map-nav. Full chunk payloads remain in ``demo_document_chunks`` and are - loaded separately for evidence hydration. + These rows are a derived index of the leaf and interstitial units used + by classic map-unit BM25 and agent recall. Full chunk payloads remain in + ``demo_document_chunks`` and are loaded separately for evidence hydration. """ __tablename__ = "demo_document_map_units" @@ -433,6 +434,74 @@ class DemoDocumentMapUnitIndex(Base): ) +class DemoDocumentPageEmbedding(Base): + """Persisted multimodal page embedding for one shared demo document revision.""" + + __tablename__ = "demo_document_page_embeddings" + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"dpe_{uuid4().hex[:12]}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="__knowhere_demo__" + ) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("demo_documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="RESTRICT"), nullable=False + ) + chunk_id: Mapped[str] = mapped_column(String(64), nullable=False) + page_num: Mapped[int] = mapped_column(Integer, nullable=False) + modality: Mapped[str] = mapped_column( + String(16), nullable=False, default="visual" + ) + model: Mapped[str] = mapped_column( + String(64), nullable=False, default="google/gemini-embedding-2" + ) + dims: Mapped[int] = mapped_column(Integer, nullable=False, default=3072) + embedding: Mapped[Any] = mapped_column(HALFVEC(3072), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + ForeignKeyConstraint( + ["job_result_id", "document_id"], + ["job_results.id", "job_results.demo_document_id"], + ondelete="RESTRICT", + name="fk_demo_documentpageembedding_revision", + ), + UniqueConstraint( + "document_id", + "job_result_id", + "page_num", + name="uq_demo_document_page_embeddings_revision_page", + ), + Index( + "idx_demo_document_page_embeddings_scope", + "user_id", + "namespace", + "job_result_id", + ), + Index( + "idx_demo_document_page_embeddings_revision", + "document_id", + "job_result_id", + "page_num", + ), + Index( + "idx_demo_document_page_embeddings_hnsw_cosine", + "embedding", + postgresql_using="hnsw", + postgresql_ops={"embedding": "halfvec_cosine_ops"}, + ), + ) + + class DemoRetrievalNamespaceGeneration(Base): """Monotonic serving generation for one user-owned namespace.""" @@ -508,8 +577,9 @@ class DemoRetrievalNamespaceMapSnapshot(Base): """Persisted namespace-level MAP (sections + chunk index + map units). Incrementally patched at publish/archive time (one document's subtree at - a time); query time reads this row directly instead of merging per-file - manifests. Overwritten in place -- no generation history is retained. + a time). Live query and ``corpus.outline`` read ``document_sections`` + instead of this blob. Overwritten in place -- no generation history is + retained. """ __tablename__ = "demo_retrieval_namespace_map_snapshots" diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index bd873b2c9..79acbe28a 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional from uuid import uuid4 +from pgvector.sqlalchemy import HALFVEC from sqlalchemy import ( JSON, DateTime, @@ -141,6 +142,46 @@ class DocumentSection(Base): ) +class DocumentScopeEnrichment(Base): + """Per-scope Phase B enrichment state for one published page-memory revision.""" + + __tablename__ = "document_scope_enrichments" + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"dse_{uuid4().hex[:12]}" + ) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + scope_id: Mapped[str] = mapped_column(String(64), nullable=False) + page_start: Mapped[int] = mapped_column(Integer, nullable=False) + page_end: Mapped[int] = mapped_column(Integer, nullable=False) + state: Mapped[str] = mapped_column(String(32), nullable=False, default="coarse") + priority: Mapped[Optional[str]] = mapped_column(String(32), nullable=True) + attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + last_hit_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True) + refined_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True) + updated_at: Mapped[datetime] = mapped_column( + DateTime, + default=utc_now_naive, + onupdate=utc_now_naive, + nullable=False, + ) + + __table_args__ = ( + UniqueConstraint( + "job_result_id", + "scope_id", + name="uq_document_scope_enrichments_revision_scope", + ), + ) + + class DocumentChunk(Base): """Canonical retrieval payload row for one published document revision.""" @@ -221,9 +262,9 @@ class DocumentChunk(Base): class DocumentMapUnit(Base): """Persisted lexical map unit for one document revision. - These rows are a derived index of the exact leaf and interstitial units - used by map-nav. Full chunk payloads remain in ``document_chunks`` and are - loaded separately for evidence hydration. + These rows are a derived index of the leaf and interstitial units used + by classic map-unit BM25 and agent recall. Full chunk payloads remain in + ``document_chunks`` and are loaded separately for evidence hydration. """ __tablename__ = "document_map_units" @@ -360,6 +401,68 @@ class DocumentMapUnitIndex(Base): ) +class DocumentPageEmbedding(Base): + """Persisted multimodal page embedding for one published document revision.""" + + __tablename__ = "document_page_embeddings" + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: f"dpe_{uuid4().hex[:12]}" + ) + user_id: Mapped[str] = mapped_column(Text, nullable=False) + namespace: Mapped[str] = mapped_column( + String(255), nullable=False, default="default" + ) + document_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("documents.document_id", ondelete="CASCADE"), + nullable=False, + ) + job_result_id: Mapped[str] = mapped_column( + String(36), ForeignKey("job_results.id", ondelete="CASCADE"), nullable=False + ) + chunk_id: Mapped[str] = mapped_column(String(64), nullable=False) + page_num: Mapped[int] = mapped_column(Integer, nullable=False) + modality: Mapped[str] = mapped_column( + String(16), nullable=False, default="visual" + ) + model: Mapped[str] = mapped_column( + String(64), nullable=False, default="google/gemini-embedding-2" + ) + dims: Mapped[int] = mapped_column(Integer, nullable=False, default=3072) + embedding: Mapped[Any] = mapped_column(HALFVEC(3072), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "document_id", + "job_result_id", + "page_num", + name="uq_document_page_embeddings_revision_page", + ), + Index( + "idx_document_page_embeddings_scope", + "user_id", + "namespace", + "job_result_id", + ), + Index( + "idx_document_page_embeddings_revision", + "document_id", + "job_result_id", + "page_num", + ), + Index( + "idx_document_page_embeddings_hnsw_cosine", + "embedding", + postgresql_using="hnsw", + postgresql_ops={"embedding": "halfvec_cosine_ops"}, + ), + ) + + class RetrievalNamespaceGeneration(Base): """Monotonic serving generation for one user-owned namespace.""" @@ -427,8 +530,9 @@ class RetrievalNamespaceMapSnapshot(Base): """Persisted namespace-level MAP (sections + chunk index + map units). Incrementally patched at publish/archive time (one document's subtree at - a time); query time reads this row directly instead of merging per-file - manifests. Overwritten in place -- no generation history is retained. + a time). Live query and ``corpus.outline`` read ``document_sections`` + instead of this blob. Overwritten in place -- no generation history is + retained. """ __tablename__ = "retrieval_namespace_map_snapshots" diff --git a/packages/shared-python/shared/models/schemas/job.py b/packages/shared-python/shared/models/schemas/job.py index e16bb65ff..87471ff90 100644 --- a/packages/shared-python/shared/models/schemas/job.py +++ b/packages/shared-python/shared/models/schemas/job.py @@ -7,6 +7,8 @@ from shared.models.schemas.llm_config import LLMConfig +IngestionMode = Literal["full", "progressive"] + class WebhookConfig(BaseModel): """Webhook configuration.""" @@ -95,6 +97,16 @@ class JobCreate(JobCreateBase): class JobCreateV2(JobCreateBase): """Public v2 request payload for creating a job.""" + ingestion_mode: IngestionMode = Field( + "full", + description=( + "Ingestion mode for page-memory documents. 'full' runs the " + "complete C1-C7 pipeline synchronously before publishing; " + "'progressive' publishes an askable per-page text baseline first, " + "then structures coarse scopes in the background and refines hit " + "scopes on demand." + ), + ) llm_config: Optional[LLMConfig] = Field( None, description=( diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py index 9b5c7d930..76082ecb1 100644 --- a/packages/shared-python/shared/models/schemas/job_metadata.py +++ b/packages/shared-python/shared/models/schemas/job_metadata.py @@ -33,6 +33,9 @@ class JobMetadataBase(BaseModel): None, description="Resolved internal ingestion generation" ) parse_track: Optional[str] = Field(None, description="Resolved parser track") + ingestion_mode: Optional[str] = Field( + None, description="Resolved ingestion mode (full | progressive)" + ) page_memory_config: Optional[Dict[str, Any]] = Field( None, description="Resolved page-memory worker configuration" ) @@ -77,6 +80,7 @@ def create_from_request( parse_track: str = "chunk", processing_generation: str = "legacy_chunk", page_memory_config: PageMemoryConfig | Dict[str, Any] | None = None, + ingestion_mode: str | None = None, **kwargs, ) -> Dict[str, Any]: """Build metadata from a public job request.""" @@ -86,6 +90,11 @@ def create_from_request( resolved_page_memory_config = page_memory_config.to_dict() else: resolved_page_memory_config = page_memory_config + resolved_ingestion_mode = _normalize_ingestion_mode( + ingestion_mode + if ingestion_mode is not None + else getattr(request, "ingestion_mode", None) + ) # v2-only field; getattr keeps v1 JobCreate (no llm_config) safe. raw_llm_config = getattr(request, "llm_config", None) llm_config_payload: Dict[str, Any] | None = None @@ -99,6 +108,7 @@ def create_from_request( "namespace": namespace, "document_id": request.document_id, "parse_track": parse_track, + "ingestion_mode": resolved_ingestion_mode, "processing_generation": processing_generation, "parsing_params": ( request.parsing_params.model_dump() if request.parsing_params else None @@ -193,6 +203,13 @@ def get_parse_track(metadata: Optional[Dict[str, Any]]) -> str: or "chunk" ) + @staticmethod + def get_ingestion_mode(metadata: Optional[Dict[str, Any]]) -> str: + """Return the ingestion mode stored in metadata (defaults to 'full').""" + return _normalize_ingestion_mode( + JobMetadataHelper.get_string_field(metadata, "ingestion_mode", "full") + ) + @staticmethod def get_api_version(metadata: Optional[Dict[str, Any]]) -> str: """Return the API version stored in metadata.""" @@ -274,6 +291,13 @@ def get_llm_config(metadata: Optional[Dict[str, Any]]) -> LLMConfig | None: return parse_llm_config(raw) +def _normalize_ingestion_mode(raw: object) -> str: + value = str(raw or "").strip().lower() + if value == "progressive": + return "progressive" + return "full" + + def _mask_llm_config_in_request(payload: Dict[str, Any]) -> Dict[str, Any]: """Redact api_key values inside llm_config for public request snapshots.""" llm_config = payload.get("llm_config") diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index 5ef96b8d1..fb0d83b6f 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -25,14 +25,13 @@ class PageMemoryConfig: tag_concurrency: int = 4 title_detection_concurrency: int = 3 node_assembly_concurrency: int = 3 - fine_min_pages: int = 4 + fine_min_pages: int = 5 hierarchy_model: str | None = None hierarchy_max_tokens: int = 2000 max_heading_depth: int = 6 asset_extraction_enabled: bool = True asset_summary_enabled: bool = False asset_model: str = "" - asset_max_pages: int | None = None asset_confidence_threshold: float = 0.3 asset_summary_concurrency: int = 4 table_engine: str = "tabula" @@ -109,7 +108,6 @@ def from_mapping(cls, value: object) -> Self: default.asset_summary_enabled, ), asset_model=_as_str(value.get("asset_model"), default.asset_model), - asset_max_pages=_as_optional_int(value.get("asset_max_pages")), asset_confidence_threshold=_as_float( value.get("asset_confidence_threshold"), default.asset_confidence_threshold, @@ -144,16 +142,6 @@ def _as_int(value: object, default: int) -> int: return resolved if resolved > 0 else default -def _as_optional_int(value: object) -> int | None: - if value in (None, ""): - return None - try: - resolved = int(str(value)) - except (TypeError, ValueError): - return None - return resolved if resolved > 0 else None - - def _as_float(value: object, default: float) -> float: try: resolved = float(str(value)) diff --git a/packages/shared-python/shared/services/ai/embedding/__init__.py b/packages/shared-python/shared/services/ai/embedding/__init__.py new file mode 100644 index 000000000..32c1c5a58 --- /dev/null +++ b/packages/shared-python/shared/services/ai/embedding/__init__.py @@ -0,0 +1,17 @@ +"""Multimodal embedding client package.""" + +from shared.services.ai.embedding.multimodal_client import ( + DEFAULT_EMBEDDING_DIM, + DEFAULT_EMBEDDING_MODEL, + embed_images, + embed_query, + resolve_embedding_model, +) + +__all__ = [ + "DEFAULT_EMBEDDING_DIM", + "DEFAULT_EMBEDDING_MODEL", + "embed_images", + "embed_query", + "resolve_embedding_model", +] diff --git a/packages/shared-python/shared/services/ai/embedding/multimodal_client.py b/packages/shared-python/shared/services/ai/embedding/multimodal_client.py new file mode 100644 index 000000000..a36129b3c --- /dev/null +++ b/packages/shared-python/shared/services/ai/embedding/multimodal_client.py @@ -0,0 +1,357 @@ +"""Vercel AI Gateway client for Gemini Embedding 2 page-image and query vectors. + +Page images and retrieval queries share one 3072-dimensional space. Images must +be sent as Gemini multimodal content; a data-URI string in ``input`` is treated +as text by the gateway and is not a visual embedding. +""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import math +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any, cast + +import httpx +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import LLMServiceException +from shared.services.ai.token_tracking import ( + bind_token_tracker, + get_current_token_tracker_root_id, + record_tokens, +) +from shared.services.http.client_pool import get_sync_client +from shared.utils.security_utils import mask_api_key + +DEFAULT_EMBEDDING_MODEL = "google/gemini-embedding-2" +DEFAULT_EMBEDDING_DIM = 3072 +DEFAULT_EMBEDDING_ENDPOINT = "https://ai-gateway.vercel.sh/v1/embeddings" +_MAX_EMBEDDING_RETRIES = 3 +_RETRY_BACKOFF_CAP_SECONDS = 8.0 + + +def resolve_embedding_model(model: str | None = None) -> str: + """Resolve the active multimodal embedding model.""" + candidate = str( + model if model is not None else getattr(settings, "EMBEDDING_MODEL", "") + ).strip() + return candidate or DEFAULT_EMBEDDING_MODEL + + +def _resolve_embedding_dim() -> int: + raw = int(getattr(settings, "EMBEDDING_DIM", DEFAULT_EMBEDDING_DIM) or DEFAULT_EMBEDDING_DIM) + return raw if raw > 0 else DEFAULT_EMBEDDING_DIM + + +def _resolve_embedding_endpoint() -> str: + endpoint = str(getattr(settings, "EMBEDDING_ENDPOINT", "") or "").strip() + return endpoint or DEFAULT_EMBEDDING_ENDPOINT + + +def _resolve_vercel_api_key() -> str: + return str(getattr(settings, "VERCEL_KEY", "") or "").strip() + + +def _should_mock_embeddings() -> bool: + return bool(getattr(settings, "LLM_MOCK_ENABLED", False)) + + +def _build_deterministic_mock_vector(seed_bytes: bytes, *, dim: int) -> list[float]: + values: list[float] = [] + counter = 0 + while len(values) < dim: + digest = hashlib.sha256(seed_bytes + counter.to_bytes(4, "big")).digest() + for idx in range(0, len(digest), 2): + if len(values) >= dim: + break + raw = int.from_bytes(digest[idx : idx + 2], "big", signed=True) + values.append(raw / 32768.0) + counter += 1 + norm = math.sqrt(sum(val * val for val in values)) or 1.0 + return [val / norm for val in values] + + +def _read_image_payload(image_input: str | Path | bytes) -> tuple[str, bytes]: + if isinstance(image_input, bytes): + return "image/png", image_input + path = Path(image_input) + raw_bytes = path.read_bytes() + suffix = path.suffix.lower() + if suffix in {".jpg", ".jpeg"}: + return "image/jpeg", raw_bytes + if suffix == ".webp": + return "image/webp", raw_bytes + return "image/png", raw_bytes + + +def _build_embedding_request( + *, + model: str, + text: str | None = None, + image_mime: str | None = None, + image_bytes: bytes | None = None, +) -> dict[str, Any]: + if image_bytes is not None: + encoded = base64.b64encode(image_bytes).decode("ascii") + mime_type = str(image_mime or "image/png").strip() or "image/png" + return { + "model": model, + "input": "", + "providerOptions": { + "google": { + "content": [ + [{"inlineData": {"mimeType": mime_type, "data": encoded}}] + ] + } + }, + } + return {"model": model, "input": str(text or "")} + + +def _extract_embedding_vector( + payload: dict[str, Any], + *, + expected_dim: int, + model: str, +) -> list[float]: + data = payload.get("data") + first = data[0] if isinstance(data, list) and data else None + raw_vec = first.get("embedding") if isinstance(first, dict) else None + if not isinstance(raw_vec, list) or not raw_vec: + raise LLMServiceException( + internal_message=f"Vercel embedding returned no embeddings: {payload}", + provider=model, + ) + vector = [float(x) for x in raw_vec] + if len(vector) != expected_dim: + raise LLMServiceException( + internal_message=( + f"Vercel embedding dimension mismatch: " + f"expected {expected_dim}, got {len(vector)}" + ), + provider=model, + ) + return vector + + +def _extract_embedding_usage(payload: dict[str, Any]) -> dict[str, int]: + usage = payload.get("usage") + if not isinstance(usage, dict): + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + input_tokens = int(usage.get("prompt_tokens") or usage.get("total_tokens") or 0) + return { + "prompt_tokens": input_tokens, + "completion_tokens": 0, + "total_tokens": input_tokens, + } + + +def _is_retryable_http_status(status_code: int) -> bool: + return status_code == 429 or 500 <= status_code < 600 + + +def _sleep_before_retry(attempt: int) -> None: + delay = min(float(2**attempt), _RETRY_BACKOFF_CAP_SECONDS) + time.sleep(delay) + + +def _call_gateway_embedding( + *, + text: str | None = None, + image_mime: str | None = None, + image_bytes: bytes | None = None, + seed_bytes: bytes, + model: str | None = None, + usage_task: str = "embedding.multimodal", +) -> list[float]: + resolved_model = resolve_embedding_model(model) + expected_dim = _resolve_embedding_dim() + + if _should_mock_embeddings(): + record_tokens( + {"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}, + model=resolved_model, + task=usage_task, + ) + return _build_deterministic_mock_vector(seed_bytes, dim=expected_dim) + + api_key = _resolve_vercel_api_key() + if not api_key: + raise LLMServiceException( + internal_message="VERCEL_KEY is required for embeddings", + provider=resolved_model, + ) + + endpoint = _resolve_embedding_endpoint() + timeout_seconds = float(getattr(settings, "OPENAI_CLIENT_TIMEOUT", 300) or 300) + http_client = get_sync_client() + request_body = _build_embedding_request( + model=resolved_model, + text=text, + image_mime=image_mime, + image_bytes=image_bytes, + ) + + last_error: Exception | None = None + for attempt in range(_MAX_EMBEDDING_RETRIES): + try: + response = http_client.post( + endpoint, + json=request_body, + headers={ + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + }, + timeout=timeout_seconds, + ) + if _is_retryable_http_status(response.status_code): + last_error = LLMServiceException( + internal_message=( + f"Vercel embedding HTTP {response.status_code}: {response.text}" + ), + provider=resolved_model, + ) + logger.warning( + "Vercel embedding retryable HTTP {} (attempt {}/{})", + response.status_code, + attempt + 1, + _MAX_EMBEDDING_RETRIES, + ) + if attempt < _MAX_EMBEDDING_RETRIES - 1: + _sleep_before_retry(attempt) + continue + raise last_error + + if response.status_code >= 400: + raise LLMServiceException( + internal_message=( + f"Vercel embedding HTTP {response.status_code}: {response.text}" + ), + provider=resolved_model, + ) + + payload = response.json() + if not isinstance(payload, dict): + raise LLMServiceException( + internal_message="Vercel embedding returned non-JSON object", + provider=resolved_model, + ) + vector = _extract_embedding_vector( + payload, + expected_dim=expected_dim, + model=resolved_model, + ) + usage = _extract_embedding_usage(payload) + record_tokens(usage, model=resolved_model, task=usage_task) + return vector + except LLMServiceException: + raise + except (httpx.RequestError, json.JSONDecodeError) as exc: + last_error = exc + masked_key = mask_api_key(api_key) + logger.warning( + "Embedding request failed (attempt {}/{}): model={}, api_key={}, error={}", + attempt + 1, + _MAX_EMBEDDING_RETRIES, + resolved_model, + masked_key, + exc, + ) + if attempt < _MAX_EMBEDDING_RETRIES - 1: + _sleep_before_retry(attempt) + continue + raise LLMServiceException( + internal_message=( + f"Embedding failed after {_MAX_EMBEDDING_RETRIES} attempts " + f"(api_key={masked_key}): {exc}" + ), + provider=resolved_model, + original_exception=exc, + ) from exc + + raise LLMServiceException( + internal_message=( + f"Vercel embedding retry loop exited unexpectedly: {last_error}" + ), + provider=resolved_model, + original_exception=last_error, + ) + + +def embed_query( + text: str, + *, + model: str | None = None, + usage_task: str = "retrieval.query_embed", +) -> list[float]: + """Embed a query string into the shared multimodal vector space.""" + normalized = str(text or "").strip() + if not normalized: + raise ValueError("embed_query requires non-empty query text") + return _call_gateway_embedding( + text=normalized, + seed_bytes=normalized.encode("utf-8"), + model=model, + usage_task=usage_task, + ) + + +def embed_images( + image_inputs: Sequence[str | Path | bytes], + *, + model: str | None = None, + concurrency: int | None = None, + usage_task: str = "page_memory.m0_embed", +) -> list[list[float] | None]: + """Embed page images concurrently. Failed pages after retries are skipped.""" + items = list(image_inputs) + if not items: + return [] + + default_concurrency = int( + getattr(settings, "PAGE_MEMORY_EMBED_CONCURRENCY", 8) or 8 + ) + resolved_concurrency = max(int(concurrency or default_concurrency), 1) + tracker_root_id = get_current_token_tracker_root_id() + + def _embed_one(index: int) -> tuple[int, list[float] | None]: + with bind_token_tracker(tracker_root_id): + try: + mime_type, raw_bytes = _read_image_payload(items[index]) + vec = _call_gateway_embedding( + image_mime=mime_type, + image_bytes=raw_bytes, + seed_bytes=raw_bytes, + model=model, + usage_task=usage_task, + ) + return index, vec + except Exception as exc: + logger.warning( + "Skipping page image embedding after retries: index={}, error={}", + index, + exc, + ) + return index, None + + if len(items) == 1 or resolved_concurrency == 1: + return [_embed_one(idx)[1] for idx in range(len(items))] + + import gevent + from gevent.pool import Pool as GeventPool + + pool = GeventPool(size=min(resolved_concurrency, len(items))) + greenlets = [pool.spawn(_embed_one, idx) for idx in range(len(items))] + gevent.joinall(greenlets, raise_error=False) + pairs = [ + cast(tuple[int, list[float] | None], greenlet.value) + for greenlet in greenlets + ] + pairs.sort(key=lambda item: item[0]) + return [vec for _, vec in pairs] diff --git a/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py index e990ec5a4..200611a15 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py @@ -6,6 +6,9 @@ RetrievalCacheInvalidation, SyncJobPublicationFinalizer, ) +from shared.services.jobs.lifecycle.scope_enrichments import ( + enqueue_page_memory_structure_after_commit, +) from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox @@ -15,6 +18,7 @@ class PostCommitEffectPlan: default_factory=tuple ) webhook_event_ids: tuple[str, ...] = field(default_factory=tuple) + structure_job_result_ids: tuple[str, ...] = field(default_factory=tuple) @classmethod def none(cls) -> PostCommitEffectPlan: @@ -26,12 +30,18 @@ def from_success( *, cache_invalidation: RetrievalCacheInvalidation | None, webhook_event_id: str | None, + structure_job_result_id: str | None = None, ) -> PostCommitEffectPlan: return cls( retrieval_cache_invalidations=( (cache_invalidation,) if cache_invalidation else () ), webhook_event_ids=((webhook_event_id,) if webhook_event_id else ()), + structure_job_result_ids=( + (structure_job_result_id,) + if structure_job_result_id + else () + ), ) @classmethod @@ -60,5 +70,7 @@ def run(self, plan: PostCommitEffectPlan) -> None: self._publication_finalizer.invalidate_cache_after_commit( cache_invalidation ) + for job_result_id in plan.structure_job_result_ids: + enqueue_page_memory_structure_after_commit(job_result_id) for webhook_event_id in plan.webhook_event_ids: self._webhook_outbox.enqueue_event_id_after_commit(webhook_event_id) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py index af373d64a..1561b7f10 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py @@ -27,21 +27,49 @@ def upsert_job_result( result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) existing = result.scalar_one_or_none() + job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + metadata = job.job_metadata or {} if job is not None else {} + if metadata.get("corpus_target") == "DEMO": + doc_metadata = { + key: metadata[key] + for key in ( + "source_size_bytes", + "demo_asset_manifest", + "result_raw_prefix", + "parse_track", + "ingestion_mode", + "enrichment_status", + ) + if key in metadata + } + else: + doc_metadata = { + key: metadata[key] + for key in ( + "result_raw_prefix", + "parse_track", + "ingestion_mode", + "enrichment_status", + ) + if key in metadata + } + if existing: existing.delivery_mode = delivery_mode existing.inline_payload = inline_payload existing.result_s3_key = result_s3_key existing.result_size = result_size + existing.document_metadata = { + **dict(existing.document_metadata or {}), + **doc_metadata, + } db.flush() return existing - job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() - metadata = job.job_metadata or {} if job is not None else {} - demoMetadata = {key: metadata[key] for key in ("source_size_bytes", "demo_asset_manifest", "result_raw_prefix", "parse_track") if key in metadata} if metadata.get("corpus_target") == "DEMO" else {} job_result = JobResult( job_id=job_id, delivery_mode=delivery_mode, - document_metadata=demoMetadata, + document_metadata=doc_metadata, inline_payload=inline_payload, result_s3_key=result_s3_key, result_size=result_size, diff --git a/packages/shared-python/shared/services/jobs/lifecycle/scope_enrichments.py b/packages/shared-python/shared/services/jobs/lifecycle/scope_enrichments.py new file mode 100644 index 000000000..5c15ae85f --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/scope_enrichments.py @@ -0,0 +1,507 @@ +from __future__ import annotations + +import re +from datetime import datetime, timedelta +from typing import Any +from uuid import uuid4 + +from loguru import logger +from sqlalchemy import ColumnElement, and_, delete, func, or_, select, update +from sqlalchemy.orm import Session + +from shared.core.config import app_config +from shared.core.database_sync import get_sync_db_context +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentScopeEnrichment, +) +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.services.retrieval.hydration.row_utils import ( + normalize_chunk_type, + page_enrichment_state, +) +from shared.utils.utc_now import utc_now_naive + +_SCOPE_ID_PAGE_RANGE = re.compile(r"^p([1-9]\d*)-([1-9]\d*)$") +_COARSE_STATE = "coarse" +_QUEUED_STATE = "queued" +_RUNNING_STATE = "running" +_FAILED_STATE = "failed" +_RESETTABLE_STATES = (_COARSE_STATE, _QUEUED_STATE) +MAX_SCOPE_ENRICHMENT_ATTEMPTS = 3 + +ENRICH_SCOPE_TASK_NAME = ( + "app.core.tasks.enrichment_tasks.enrich_page_memory_scope" +) +BUILD_STRUCTURE_TASK_NAME = ( + "app.core.tasks.enrichment_tasks.build_page_memory_structure" +) +LOW_ENRICHMENT_QUEUE = "document_ingestion_low" +NORMAL_ENRICHMENT_QUEUE = "document_ingestion_medium" +HIGH_ENRICHMENT_QUEUE = "document_ingestion_high" +HIGH_ENRICHMENT_TASK_PRIORITY = 9 + + +def scope_claimable_filter(now: datetime) -> ColumnElement[bool]: + """Single source of truth for which enrichment rows may be claimed or re-queued. + + A row is claimable when it is still ``coarse``/``queued``, when it failed + with fewer than ``MAX_SCOPE_ENRICHMENT_ATTEMPTS`` attempts, or when it has + been stuck in ``running`` past the Celery hard time limit (indicating the + previous worker was killed by OOM/SIGKILL without releasing the row). + ``TASK_TIME_LIMIT_SECONDS`` is strictly smaller than the broker redelivery + window, so a running row older than it cannot still be owned by a live task. + """ + stale_cutoff = now - timedelta(seconds=app_config.TASK_TIME_LIMIT_SECONDS) + return or_( + DocumentScopeEnrichment.state.in_(_RESETTABLE_STATES), + and_( + DocumentScopeEnrichment.state == _FAILED_STATE, + DocumentScopeEnrichment.attempts < MAX_SCOPE_ENRICHMENT_ATTEMPTS, + ), + and_( + DocumentScopeEnrichment.state == _RUNNING_STATE, + DocumentScopeEnrichment.updated_at < stale_cutoff, + ), + ) + + +def write_coarse_scope_enrichments( + db: Session, + *, + document_id: str, + job_result_id: str, +) -> bool: + """Replace this revision's enrichment rows from published page-chunk scopes.""" + published = ( + db.execute( + select(DocumentChunk.chunk_metadata) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_type == "page") + .order_by(DocumentChunk.sort_order, DocumentChunk.id) + ) + .scalars() + .all() + ) + scopes: dict[str, tuple[int, int]] = {} + for metadata in published: + if not isinstance(metadata, dict): + continue + scope_id = str(metadata.get("scope_id") or "").strip() + if not scope_id or scope_id in scopes: + continue + page_range = _page_range_from_scope_id(scope_id) + if page_range is None: + continue + scopes[scope_id] = page_range + if not scopes: + return False + + db.execute( + delete(DocumentScopeEnrichment).where( + DocumentScopeEnrichment.job_result_id == job_result_id + ) + ) + now = utc_now_naive() + db.add_all( + [ + DocumentScopeEnrichment( + id=f"dse_{uuid4().hex[:12]}", + document_id=document_id, + job_result_id=job_result_id, + scope_id=scope_id, + page_start=page_start, + page_end=page_end, + state=_COARSE_STATE, + attempts=0, + updated_at=now, + ) + for scope_id, (page_start, page_end) in scopes.items() + ] + ) + db.flush() + return True + + +def enqueue_low_priority_scope_enrichments_after_commit( + job_result_id: str, +) -> list[str]: + """Mark all claimable scopes of a revision ``queued(low)`` and enqueue largest first.""" + if not job_result_id: + return [] + + enqueued_scopes: list[str] = [] + now = utc_now_naive() + try: + with get_sync_db_context() as db: + candidates = ( + db.execute( + select( + DocumentScopeEnrichment.scope_id, + DocumentScopeEnrichment.page_start, + DocumentScopeEnrichment.page_end, + ) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(scope_claimable_filter(now)) + .order_by( + ( + DocumentScopeEnrichment.page_end + - DocumentScopeEnrichment.page_start + ).desc(), + DocumentScopeEnrichment.page_start.asc(), + DocumentScopeEnrichment.scope_id.asc(), + ) + ) + .all() + ) + for scope_id, _page_start, _page_end in candidates: + claimed = db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == str(scope_id)) + .where(scope_claimable_filter(now)) + .values( + state=_QUEUED_STATE, + priority="low", + updated_at=now, + ) + .returning(DocumentScopeEnrichment.scope_id) + ).scalar_one_or_none() + if claimed is not None: + enqueued_scopes.append(str(claimed)) + except Exception as exc: + logger.warning( + "Failed to transition scope enrichments to queued after commit " + f"(ignored): job_result_id={job_result_id}, error={exc}" + ) + return [] + + if not enqueued_scopes: + return [] + + sent_scopes: list[str] = [] + try: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + for scope_id in enqueued_scopes: + celery_app.send_task( + ENRICH_SCOPE_TASK_NAME, + args=[job_result_id, scope_id], + queue=LOW_ENRICHMENT_QUEUE, + ) + sent_scopes.append(scope_id) + except Exception as exc: + logger.warning( + "Failed to send low-priority scope enrichment tasks after commit " + f"(reverting unsent to coarse): job_result_id={job_result_id}, error={exc}" + ) + unsent = [s for s in enqueued_scopes if s not in sent_scopes] + if unsent: + _revert_queued_scopes_to_coarse( + job_result_id=job_result_id, + scope_ids=unsent, + ) + return sent_scopes + return sent_scopes + + +def enqueue_page_memory_structure_after_commit(job_result_id: str) -> bool: + """Enqueue the background M0 -> coarse structure task at normal priority.""" + if not job_result_id: + return False + try: + from shared.core.celery_app import get_celery_app + + get_celery_app().send_task( + BUILD_STRUCTURE_TASK_NAME, + args=[job_result_id], + queue=NORMAL_ENRICHMENT_QUEUE, + ) + return True + except Exception as exc: + logger.warning( + "Failed to send build_page_memory_structure task after commit " + f"(ignored): job_result_id={job_result_id}, error={exc}" + ) + return False + + +def _revert_queued_scopes_to_coarse( + *, + job_result_id: str, + scope_ids: list[str], +) -> None: + now = utc_now_naive() + try: + with get_sync_db_context() as db: + db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id.in_(scope_ids)) + .where(DocumentScopeEnrichment.state == _QUEUED_STATE) + .values(state=_COARSE_STATE, priority=None, updated_at=now) + ) + except Exception as exc: + logger.warning( + f"Failed to revert unsent scopes to coarse for {job_result_id}: {exc}" + ) + + +async def promote_scopes_on_read( + db: Any, + rows: list[dict[str, Any]], +) -> list[tuple[str, str]]: + """Promote coarse/queued page scopes hit by ``corpus.read`` to high priority.""" + targets: list[tuple[str, str]] = [] + seen: set[tuple[str, str]] = set() + for row in rows: + if normalize_chunk_type(row.get("chunk_type")) != "page": + continue + if page_enrichment_state(row) not in (_COARSE_STATE, _QUEUED_STATE): + continue + metadata = row.get("chunk_metadata") or row.get("metadata") or {} + scope_id = str(metadata.get("scope_id") or "").strip() + job_result_id = str(row.get("job_result_id") or "").strip() + if not scope_id or not job_result_id: + continue + key = (job_result_id, scope_id) + if key in seen: + continue + seen.add(key) + targets.append(key) + + if not targets: + return [] + + promoted: list[tuple[str, str]] = [] + now = utc_now_naive() + try: + for job_result_id, scope_id in targets: + result = await db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .where(scope_claimable_filter(now)) + .where( + or_( + DocumentScopeEnrichment.state != _QUEUED_STATE, + DocumentScopeEnrichment.priority.is_(None), + DocumentScopeEnrichment.priority != "high", + ) + ) + .values( + state=_QUEUED_STATE, + priority="high", + last_hit_at=now, + updated_at=now, + ) + .returning(DocumentScopeEnrichment.scope_id) + ) + if result.scalar_one_or_none() is not None: + promoted.append((job_result_id, scope_id)) + if promoted: + await db.commit() + except Exception as exc: + logger.warning( + f"Failed to promote scope enrichment on corpus.read (ignored): {exc}" + ) + return [] + + if not promoted: + return [] + + sent: list[tuple[str, str]] = [] + try: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + for job_result_id, scope_id in promoted: + celery_app.send_task( + ENRICH_SCOPE_TASK_NAME, + args=[job_result_id, scope_id], + queue=HIGH_ENRICHMENT_QUEUE, + priority=HIGH_ENRICHMENT_TASK_PRIORITY, + ) + sent.append((job_result_id, scope_id)) + except Exception as exc: + logger.warning( + f"Failed to enqueue promoted scope enrichment on corpus.read (reverting unsent to coarse): {exc}" + ) + now_revert = utc_now_naive() + unsent_pairs = [item for item in promoted if item not in sent] + try: + for job_result_id, scope_id in unsent_pairs: + await db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == job_result_id) + .where(DocumentScopeEnrichment.scope_id == scope_id) + .where(DocumentScopeEnrichment.state == _QUEUED_STATE) + .values(state=_COARSE_STATE, priority=None, updated_at=now_revert) + ) + if unsent_pairs: + await db.commit() + except Exception as revert_exc: + logger.warning( + f"Failed to revert unsent promoted scopes to coarse: {revert_exc}" + ) + return sent + return sent + + +def _structure_claim_still_active( + *, + doc_meta: object, + job_meta: object, + now: datetime, +) -> bool: + doc_dict = doc_meta if isinstance(doc_meta, dict) else {} + job_dict = job_meta if isinstance(job_meta, dict) else {} + raw_claimed_at = doc_dict.get("structure_claimed_at") or job_dict.get("structure_claimed_at") + if not raw_claimed_at: + return False + try: + claimed_at = datetime.fromisoformat( + str(raw_claimed_at).replace("Z", "+00:00") + ).replace(tzinfo=None) + return (now - claimed_at) < timedelta(seconds=app_config.TASK_TIME_LIMIT_SECONDS) + except Exception: + return False + + +def sweep_progressive_enrichments_sync( + *, + structure_timeout_seconds: int = 300, +) -> dict[str, Any]: + """Re-enqueue revisions stuck at ``text`` and retry failed/un-enqueued scopes.""" + now = utc_now_naive() + structure_cutoff = now - timedelta(seconds=structure_timeout_seconds) + retried_scopes: list[tuple[str, str]] = [] + + with get_sync_db_context() as db: + structure_rows = db.execute( + select( + JobResult.id, + JobResult.document_metadata, + Job.job_metadata, + ) + .select_from(JobResult) + .join( + Document, + and_( + Document.current_job_result_id == JobResult.id, + Document.status == "active", + ), + ) + .join(Job, Job.job_id == JobResult.job_id) + .where(Job.updated_at < structure_cutoff) + .where( + func.coalesce( + JobResult.document_metadata["enrichment_status"].as_string(), + Job.job_metadata["enrichment_status"].as_string(), + ) + == "text" + ) + ).all() + reenqueued_structures = [ + str(revision_id) + for revision_id, doc_meta, job_meta in structure_rows + if not _structure_claim_still_active( + doc_meta=doc_meta, job_meta=job_meta, now=now + ) + ] + + scope_candidates = ( + db.execute( + select( + DocumentScopeEnrichment.job_result_id, + DocumentScopeEnrichment.scope_id, + ) + .select_from(DocumentScopeEnrichment) + .join( + Document, + and_( + Document.document_id == DocumentScopeEnrichment.document_id, + Document.current_job_result_id + == DocumentScopeEnrichment.job_result_id, + Document.status == "active", + ), + ) + .where( + or_( + DocumentScopeEnrichment.state == _COARSE_STATE, + and_( + DocumentScopeEnrichment.state == _FAILED_STATE, + DocumentScopeEnrichment.attempts + < MAX_SCOPE_ENRICHMENT_ATTEMPTS, + ), + and_( + DocumentScopeEnrichment.state == _RUNNING_STATE, + DocumentScopeEnrichment.updated_at + < now - timedelta(seconds=app_config.TASK_TIME_LIMIT_SECONDS), + ), + ) + ) + ) + .all() + ) + for rev_id, scope_id in scope_candidates: + claimed = db.execute( + update(DocumentScopeEnrichment) + .where(DocumentScopeEnrichment.job_result_id == str(rev_id)) + .where(DocumentScopeEnrichment.scope_id == str(scope_id)) + .where(scope_claimable_filter(now)) + .values(state=_QUEUED_STATE, priority="low", updated_at=now) + .returning(DocumentScopeEnrichment.scope_id) + ).scalar_one_or_none() + if claimed is not None: + retried_scopes.append((str(rev_id), str(claimed))) + + sent_structures: list[str] = [] + for rev_id in reenqueued_structures: + if enqueue_page_memory_structure_after_commit(rev_id): + sent_structures.append(rev_id) + + sent_scopes: list[tuple[str, str]] = [] + if retried_scopes: + try: + from shared.core.celery_app import get_celery_app + + celery_app = get_celery_app() + for rev_id, scope_id in retried_scopes: + celery_app.send_task( + ENRICH_SCOPE_TASK_NAME, + args=[rev_id, scope_id], + queue=LOW_ENRICHMENT_QUEUE, + ) + sent_scopes.append((rev_id, scope_id)) + except Exception as exc: + logger.warning(f"Sweeper failed to send scope enrichment tasks: {exc}") + by_rev: dict[str, list[str]] = {} + for rev_id, scope_id in retried_scopes: + if (rev_id, scope_id) not in sent_scopes: + by_rev.setdefault(rev_id, []).append(scope_id) + for rev_id, unsent in by_rev.items(): + _revert_queued_scopes_to_coarse( + job_result_id=rev_id, scope_ids=unsent + ) + + return { + "status": "success", + "reenqueued_structures": sent_structures, + "retried_scopes": sent_scopes, + } + + +def _page_range_from_scope_id(scope_id: str) -> tuple[int, int] | None: + matched = _SCOPE_ID_PAGE_RANGE.fullmatch(scope_id) + if matched is None: + return None + page_start = int(matched.group(1)) + page_end = int(matched.group(2)) + if page_end < page_start: + return None + return page_start, page_end diff --git a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py index e6e0d332d..c01e909b7 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py @@ -5,9 +5,12 @@ from typing import Any, Literal, cast from loguru import logger +from sqlalchemy import select from sqlalchemy.orm import Session from shared.core.state_machine.service_sync import SyncStateMachineService +from shared.models.database.job import Job +from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.jobs.lifecycle.post_commit_effects import PostCommitEffectPlan from shared.services.jobs.lifecycle.publication import SyncJobPublicationFinalizer from shared.services.jobs.lifecycle.publication_trace import PublicationTrace @@ -117,6 +120,22 @@ def finalize( document_top_summary=safe_document_top_summary, trace=trace, ) + published_document_id = ( + publication_outcome.published_document_state.document_id + if publication_outcome.published_document_state is not None + else None + ) + job_row = db.execute( + select(Job.job_metadata).where(Job.job_id == job_id) + ).scalar_one_or_none() + job_metadata = job_row if isinstance(job_row, dict) else {} + should_enqueue_structure = bool( + published_document_id + and JobMetadataHelper.get_parse_track(job_metadata) == "page_memory" + and JobMetadataHelper.get_ingestion_mode(job_metadata) == "progressive" + and str(job_metadata.get("enrichment_status") or "").strip().lower() + == "text" + ) transition_outcome = self._state_machine.mark_completed_outcome( db, @@ -150,5 +169,8 @@ def finalize( post_commit_effects=PostCommitEffectPlan.from_success( cache_invalidation=publication_outcome.cache_invalidation, webhook_event_id=webhook_event.event_id if webhook_event else None, + structure_job_result_id=( + str(job_result.id) if should_enqueue_structure else None + ), ), ) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py b/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py index 206ed702c..0b133d43b 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py @@ -2,8 +2,7 @@ Agentic path when ``use_agentic`` is true, or when it is unset and the query router selects agent. Runs the ``agent_tools`` corpus registry through an LLM -tool-calling loop. Map-nav is archived under ``deprecated/mapnav/`` and is -not a live route. +tool-calling loop. Which provider runs that loop (Cursor SDK or OpenAI-compatible) is the ``AGENT_EXPLORE_HARNESS`` switch resolved via ``resolve_harness()``. diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/config.py b/packages/shared-python/shared/services/retrieval/agent_explore/config.py index 4cb318acb..7ebc06f53 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/config.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/config.py @@ -74,5 +74,5 @@ PICK_TOOL_DESCRIPTION = ( "Pick which of your latest corpus.read / corpus.outline results join the " - "evidence pool. Body text you read but do not pick cannot be read again." + "evidence pool. Body text you read cannot be read again." ) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py index 1db56ac54..39a87af5a 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py @@ -18,8 +18,7 @@ step. Tokens are only summed for the episode total. Tool calls within one turn are dispatched sequentially through -``dispatch.dispatch_tool_call`` (fresh DB session per call). No import from -archived map-nav modules. +``dispatch.dispatch_tool_call`` (fresh DB session per call). """ from __future__ import annotations diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py b/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py index c84f6c17b..8d523810a 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py @@ -28,7 +28,7 @@ - After a turn where corpus.read or corpus.outline succeeds, the next turn is a pick phase: you can only call corpus.pick. Pass the ids worth keeping, or an empty list. Picked ids join the evidence pool. -- Body text you read but did not pick cannot be read again. +- Body text you read (picked or not) cannot be read again. - Each turn you see: a trace of your past calls (tool, arguments, ok or why it failed), the full results of your latest calls, the evidence pool, and the steps used so far. @@ -36,9 +36,6 @@ Call finish when the pool is enough to answer, or when you are sure the corpus has no answer. finish ends the exploration at once. -If a call needs another call's result, wait for that result instead of -issuing both in the same turn. - Tool names on the wire use "_" instead of "." (corpus_read = corpus.read). """ diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md index 09fd4de50..b9007bee7 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md +++ b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md @@ -25,3 +25,14 @@ doc_7f3a Hypertension_Guideline.pdf Addresses from this example: - {document_id: "doc_7f3a", section_path: "2 Treatment / 2.2 Drugs"} - {document_id: "doc_7f3a", chunk_id: "tbl_12"} + +Documents may still be loading progressively (`status` is `text`, `coarse`, or +`complete`): +- Pages are always searchable immediately via `corpus.recall` (lexical and + visual semantic for scanned pages), `corpus.grep` (once text/OCR is present), + and `corpus.read` (returns page text or the rendered page image). +- `corpus.outline` and `corpus.node_filter` require `coarse` or `complete`. +- Extracted figure and table tools require per-scope refinement. +A "not ready yet" tool note means background processing is still running — do +not retry that tool on the same scope in this session; answer from page text +or page images returned by `corpus.read`. diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/enrichment_state.py b/packages/shared-python/shared/services/retrieval/agent_tools/enrichment_state.py new file mode 100644 index 000000000..35c15617f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/enrichment_state.py @@ -0,0 +1,177 @@ +"""Progressive-enrichment state queries for ``corpus.*`` tools.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.job import Job +from shared.models.database.job_result import JobResult +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.hydration.row_utils import ( + extract_page_nums, + page_enrichment_state, +) + +STATUS_TEXT = "text" +STATUS_COARSE = "coarse" +STATUS_COMPLETE = "complete" +UNREFINED_PAGE_STATES = frozenset({"text", "coarse", "queued"}) + + +@dataclass(frozen=True) +class ScopePageStateSummary: + states: frozenset[str] + page_min: int | None + page_max: int | None + page_rows: tuple[dict[str, Any], ...] + + @property + def has_unrefined_pages(self) -> bool: + return bool(self.states & UNREFINED_PAGE_STATES) + + @property + def page_range_label(self) -> str: + if self.page_min is not None and self.page_max is not None: + if self.page_min == self.page_max: + return f"page {self.page_min}" + return f"pages {self.page_min}-{self.page_max}" + return "the requested scope" + + +async def load_revision_enrichment_status( + db: AsyncSession, + job_result_ids: list[str] | set[str], +) -> dict[str, str]: + """Return ``job_result_id -> status`` ('text' | 'coarse' | 'complete'). + + Absent keys on full-mode or legacy revisions resolve to ``'complete'``. + """ + ids = sorted({str(item).strip() for item in job_result_ids if str(item).strip()}) + if not ids: + return {} + out: dict[str, str] = {rev_id: STATUS_COMPLETE for rev_id in ids} + rows = ( + await db.execute( + select( + JobResult.id, + JobResult.document_metadata, + Job.job_metadata, + ) + .select_from(JobResult) + .outerjoin(Job, Job.job_id == JobResult.job_id) + .where(JobResult.id.in_(ids)) + ) + ).all() + for rev_id, doc_meta, job_meta in rows: + raw_status = None + if isinstance(doc_meta, dict): + raw_status = doc_meta.get("enrichment_status") + if not raw_status and isinstance(job_meta, dict): + raw_status = job_meta.get("enrichment_status") + normalized = str(raw_status or "").strip().lower() + if normalized in (STATUS_TEXT, STATUS_COARSE, STATUS_COMPLETE): + out[str(rev_id)] = normalized + else: + out[str(rev_id)] = STATUS_COMPLETE + return out + + +async def inspect_scope_page_states( + db: AsyncSession, + *, + namespace: str, + document_id: str, + job_result_id: str, + section_path: str | None = None, +) -> ScopePageStateSummary: + """Inspect page chunks under ``(document_id, job_result_id, section_path?)``.""" + corpus_storage = CorpusStorage.resolve_namespace(namespace) + stmt = ( + select( + corpus_storage.DocumentChunk.chunk_id, + corpus_storage.DocumentChunk.chunk_type, + corpus_storage.DocumentChunk.chunk_metadata, + corpus_storage.DocumentSection.section_path, + ) + .select_from(corpus_storage.DocumentChunk) + .outerjoin( + corpus_storage.DocumentSection, + corpus_storage.DocumentSection.section_id + == corpus_storage.DocumentChunk.section_id, + ) + .where(corpus_storage.DocumentChunk.document_id == document_id) + .where(corpus_storage.DocumentChunk.job_result_id == job_result_id) + .where(corpus_storage.DocumentChunk.chunk_type == "page") + ) + if section_path is not None: + stmt = stmt.where( + or_( + corpus_storage.DocumentSection.section_path == section_path, + corpus_storage.DocumentSection.section_path.like( + f"{section_path} / %" + ), + ) + ) + rows = (await db.execute(stmt)).all() + states: set[str] = set() + page_nums_all: list[int] = [] + page_rows: list[dict[str, Any]] = [] + for chunk_id, chunk_type, metadata, resolved_path in rows: + meta_dict = metadata if isinstance(metadata, dict) else {} + row_dict: dict[str, Any] = { + "document_id": document_id, + "job_result_id": job_result_id, + "chunk_id": str(chunk_id or ""), + "chunk_type": str(chunk_type or "page"), + "chunk_metadata": meta_dict, + "section_path": resolved_path, + } + state = page_enrichment_state(row_dict) + if state: + states.add(state) + nums = extract_page_nums(row_dict) or [] + for num in nums: + try: + page_nums_all.append(int(num)) + except (TypeError, ValueError): + continue + page_rows.append(row_dict) + return ScopePageStateSummary( + states=frozenset(states), + page_min=min(page_nums_all) if page_nums_all else None, + page_max=max(page_nums_all) if page_nums_all else None, + page_rows=tuple(page_rows), + ) + + +def format_still_structuring_note(file_names: list[str]) -> str: + joined = ", ".join(file_names) if file_names else "document" + return ( + f"{joined} is still being structured (state=text); no section outline yet. " + "Do not retry outline, node_filter, or assets on it in this session — " + "use corpus.recall (including visual semantic recall for scanned pages), " + "corpus.grep, and corpus.read on pages directly." + ) + + +def format_assets_not_refined_note(page_range_label: str, state: str = "coarse") -> str: + return ( + f"assets for {page_range_label} are not extracted yet (state={state}); " + "background refinement is queued. Do not retry corpus.assets or " + "corpus.query_table on this range in this session — use page text " + "via corpus.read, corpus.grep, or corpus.recall." + ) + + +def format_scanned_pages_pending_ocr_note(file_names: list[str]) -> str: + joined = ", ".join(file_names) if file_names else "document" + return ( + f"{joined} is a scanned document still awaiting background OCR (state=text); " + "no page text is indexed for exact string search yet. Do not retry " + "corpus.grep on it in this session — use corpus.recall (visual semantic " + "search over page images) and corpus.read to inspect page images." + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py index 06d14eb90..1bef9f6bb 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py @@ -31,11 +31,21 @@ from sqlalchemy import select +from shared.services.jobs.lifecycle.scope_enrichments import ( + promote_scopes_on_read, +) from shared.services.retrieval.agent_tools.asset_hosts import ( hosted_section_path, in_scope_hosts, load_asset_hosts, ) +from shared.services.retrieval.agent_tools.enrichment_state import ( + STATUS_TEXT, + format_assets_not_refined_note, + format_still_structuring_note, + inspect_scope_page_states, + load_revision_enrichment_status, +) from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, @@ -189,6 +199,14 @@ async def _forward_search( if len(rows) >= ctx.budget.max_items: break + if not rows: + unrefined_note = await _check_unrefined_assets_note(ctx, scope=scope) + if unrefined_note is not None: + return ToolResult( + text=unrefined_note, + payload={"rows": [], "details": {"unrefined": True}}, + refs=[], + ) lines = [f"assets={len(rows)}", *(format_row(row) for row in rows)] return ToolResult( text="\n".join(lines), @@ -197,6 +215,96 @@ async def _forward_search( ) +async def _check_unrefined_assets_note( + ctx: ToolContext, + *, + scope: list[ScopeTarget], +) -> str | None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) + targets = scope + if not targets: + docs = ( + ( + await ctx.db.execute( + select(corpusStorage.Document) + .where( + ctx.document_scope.predicate( + corpusStorage.Document.document_id + ), + corpusStorage.Document.user_id + == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.Document.namespace == ctx.namespace, + corpusStorage.Document.status == "active", + ) + ) + ) + .scalars() + .all() + ) + targets = [ + ScopeTarget(d.document_id, rev, None) + for d in docs + if ( + rev := CorpusRevisionContext.resolve_revision( + d.document_id, d.current_job_result_id + ) + ) + ] + if not targets: + return None + + revision_status = await load_revision_enrichment_status( + ctx.db, + {t.job_result_id for t in targets}, + ) + text_targets = [ + t for t in targets if revision_status.get(t.job_result_id) == STATUS_TEXT + ] + if text_targets and len(text_targets) == len(targets): + doc_rows = ( + ( + await ctx.db.execute( + select(corpusStorage.Document).where( + corpusStorage.Document.document_id.in_( + [t.document_id for t in text_targets] + ) + ) + ) + ) + .scalars() + .all() + ) + names = [d.source_file_name or d.document_id for d in doc_rows] + return format_still_structuring_note(names) + + promote_rows: list[dict[str, Any]] = [] + range_labels: list[str] = [] + seen_states: set[str] = set() + for target in targets: + summary = await inspect_scope_page_states( + ctx.db, + namespace=ctx.namespace, + document_id=target.document_id, + job_result_id=target.job_result_id, + section_path=target.section_path, + ) + if not summary.has_unrefined_pages: + continue + seen_states.update(summary.states) + range_labels.append(summary.page_range_label) + promote_rows.extend(summary.page_rows) + + if not range_labels: + return None + if promote_rows: + await promote_scopes_on_read(ctx.db, promote_rows) + state_label = "text" if "text" in seen_states else "coarse" + return format_assets_not_refined_note( + ", ".join(dict.fromkeys(range_labels)), + state=state_label, + ) + + async def _reverse_lookup( ctx: ToolContext, *, diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py index 2815cc0cc..1d407b740 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py @@ -32,6 +32,11 @@ from sqlalchemy import func, or_, select from shared.services.retrieval.agent_tools.asset_hosts import host_paths_for_hits +from shared.services.retrieval.agent_tools.enrichment_state import ( + STATUS_TEXT, + format_scanned_pages_pending_ocr_note, + load_revision_enrichment_status, +) from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, @@ -353,6 +358,19 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: _fold_line(str(group[0]["source_file_name"] or ""), document_id, len(group)) ) + if not rows: + scanned_note = await _check_scanned_pending_ocr_note(ctx, scope=scope) + if scanned_note is not None: + if lines: + lines.append(scanned_note) + else: + lines = [scanned_note] + return ToolResult( + text="\n".join(lines), + payload={"rows": [], "details": {"enrichment_status": STATUS_TEXT}}, + refs=[], + ) + return ToolResult( text="\n".join(lines), payload={"rows": rows, "details": {}}, @@ -361,3 +379,95 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: for hit in shown_hits ], ) + + +async def _check_scanned_pending_ocr_note( + ctx: ToolContext, + *, + scope: list[ScopeTarget], +) -> str | None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(ctx.namespace) + targets = scope + if not targets: + docs = ( + ( + await ctx.db.execute( + select(corpusStorage.Document).where( + ctx.document_scope.predicate( + corpusStorage.Document.document_id + ), + corpusStorage.Document.user_id + == corpusStorage.resolve_owner(ctx.user_id), + corpusStorage.Document.namespace == ctx.namespace, + corpusStorage.Document.status == "active", + ) + ) + ) + .scalars() + .all() + ) + targets = [ + ScopeTarget(d.document_id, rev, None) + for d in docs + if ( + rev := CorpusRevisionContext.resolve_revision( + d.document_id, d.current_job_result_id + ) + ) + ] + if not targets: + return None + + revision_status = await load_revision_enrichment_status( + ctx.db, + {t.job_result_id for t in targets}, + ) + text_targets = [ + t for t in targets if revision_status.get(t.job_result_id) == STATUS_TEXT + ] + if not text_targets: + return None + + scanned_doc_ids: list[str] = [] + seen_docs: set[str] = set() + for target in text_targets: + if target.document_id in seen_docs: + continue + seen_docs.add(target.document_id) + page_rows = ( + await ctx.db.execute( + select( + corpusStorage.DocumentChunk.content, + corpusStorage.DocumentChunk.chunk_metadata, + ) + .where(corpusStorage.DocumentChunk.document_id == target.document_id) + .where( + corpusStorage.DocumentChunk.job_result_id + == target.job_result_id + ) + .where(corpusStorage.DocumentChunk.chunk_type == "page") + ) + ).all() + if page_rows and all( + not str(content or "").strip() + and isinstance(meta, dict) + and bool(meta.get("page_assets")) + for content, meta in page_rows + ): + scanned_doc_ids.append(target.document_id) + + if not scanned_doc_ids: + return None + doc_rows = ( + ( + await ctx.db.execute( + select(corpusStorage.Document).where( + corpusStorage.Document.document_id.in_(scanned_doc_ids) + ) + ) + ) + .scalars() + .all() + ) + names = [d.source_file_name or d.document_id for d in doc_rows] + return format_scanned_pages_pending_ocr_note(names) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py index a312a8ac3..8c1c2543f 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py @@ -14,6 +14,10 @@ from sqlalchemy import select +from shared.services.retrieval.agent_tools.enrichment_state import ( + STATUS_COMPLETE, + load_revision_enrichment_status, +) from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, @@ -53,15 +57,30 @@ async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: .order_by(corpusStorage.Document.source_file_name) ) rows = (await ctx.db.execute(stmt)).all() + revision_ids = [ + rev + for document, _props in rows + if ( + rev := CorpusRevisionContext.resolve_revision( + document.document_id, document.current_job_result_id + ) + ) + ] + status_by_rev = await load_revision_enrichment_status(ctx.db, revision_ids) documents: list[dict[str, Any]] = [] lines: list[str] = [] for document, properties in rows: props = properties if isinstance(properties, dict) else {} + rev_id = CorpusRevisionContext.resolve_revision( + document.document_id, document.current_job_result_id + ) + status = status_by_rev.get(str(rev_id or ""), STATUS_COMPLETE) entry = { "document_id": document.document_id, - "job_result_id": CorpusRevisionContext.resolve_revision(document.document_id, document.current_job_result_id), + "job_result_id": rev_id, "source_file_name": document.source_file_name, + "status": status, "top_keywords": props.get("top_keywords") or [], "top_summary": props.get("top_summary") or "", "types": props.get("types") or {}, @@ -70,7 +89,7 @@ async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: documents.append(entry) summary_line = ( f"- {entry['source_file_name']} ({entry['document_id']}, " - f"chunks={entry['chunks_count']})" + f"status={status}, chunks={entry['chunks_count']})" ) if entry["top_summary"]: summary_line += f"\n summary: {entry['top_summary']}" diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py index 2b84d3982..4802ac353 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py @@ -2,9 +2,8 @@ Reuses the exact predicate compile/match semantics from ``scoring.node_filter_predicates`` (path/summary substring|regex, fields AND -together, terms OR together) — see that module's docstring — but walks -``document_sections`` rows for the requested scope's current revision -instead of the in-memory map-nav tree. +together, terms OR together) — see that module's docstring — and walks +``document_sections`` rows for the requested scope's current revision. Unlike a plain filter, the rendered result is not just the matched rows: it is each matched section's full ancestor chain (context) plus its entire @@ -26,6 +25,11 @@ from sqlalchemy import select +from shared.services.retrieval.agent_tools.enrichment_state import ( + STATUS_TEXT, + format_still_structuring_note, + load_revision_enrichment_status, +) from shared.services.retrieval.agent_tools.map_render import render_map from shared.services.retrieval.agent_tools.registry import ( ToolContext, @@ -156,6 +160,50 @@ async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: document_ids = scope_document_ids(targets) revision_by_doc = {target.document_id: target.job_result_id for target in targets} + revision_status = await load_revision_enrichment_status( + ctx.db, + set(revision_by_doc.values()), + ) + text_targets = [ + target + for target in targets + if revision_status.get(target.job_result_id) == STATUS_TEXT + ] + if text_targets: + documents_rows = ( + ( + await ctx.db.execute( + select(corpusStorage.Document).where( + corpusStorage.Document.document_id.in_(document_ids) + ) + ) + ) + .scalars() + .all() + ) + source_file_name_by_doc = { + d.document_id: d.source_file_name or d.document_id for d in documents_rows + } + text_names = [ + source_file_name_by_doc.get(t.document_id) or t.document_id + for t in text_targets + ] + if len(text_targets) == len(targets): + note = format_still_structuring_note(text_names) + return ToolResult( + text=note, + payload={ + "rows": [], + "details": { + "cardinality": 0, + "matched_document_ids": [], + "enrichment_status": STATUS_TEXT, + }, + }, + refs=[], + ) + else: + text_names = [] chunk_types = { str(t).strip().lower() for t in (args.get("chunk_types") or []) if str(t).strip() @@ -244,8 +292,13 @@ async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: order.append((section.section_id, parent_id)) if not matched_ids: + empty_text = ( + format_still_structuring_note(text_names) + if text_names + else "hits=0" + ) return ToolResult( - text="hits=0", + text=empty_text, payload={"rows": [], "details": {"cardinality": 0, "matched_document_ids": []}}, refs=[], ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py index b6ac97a1a..b15cc9650 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py @@ -25,6 +25,11 @@ from sqlalchemy import select +from shared.services.retrieval.agent_tools.enrichment_state import ( + STATUS_TEXT, + format_still_structuring_note, + load_revision_enrichment_status, +) from shared.services.retrieval.agent_tools.map_render import render_map from shared.services.retrieval.agent_tools.registry import ( ToolContext, @@ -102,6 +107,32 @@ async def outline(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: .all() ) source_file_name_by_doc = {d.document_id: d.source_file_name or "" for d in documents} + revision_status = await load_revision_enrichment_status( + ctx.db, + {target.job_result_id for target in targets}, + ) + text_targets = [ + target + for target in targets + if revision_status.get(target.job_result_id) == STATUS_TEXT + ] + if text_targets and len(text_targets) == len(targets): + names = [ + source_file_name_by_doc.get(t.document_id) or t.document_id + for t in text_targets + ] + note = format_still_structuring_note(names) + return ToolResult( + text=note, + payload={ + "rows": [], + "details": { + "enrichment_status": STATUS_TEXT, + "documents": source_file_name_by_doc, + }, + }, + refs=[], + ) rows_by_id: dict[str, dict[str, Any]] = {} order: list[tuple[str, str | None]] = [] @@ -200,8 +231,15 @@ def header(visible: int, hidden: int) -> str: visible_rows = [ rows_by_id[section_id] for section_id, _ in order if section_id in rendered.visible_ids ] + text_out = rendered.text + if text_targets: + names = [ + source_file_name_by_doc.get(t.document_id) or t.document_id + for t in text_targets + ] + text_out = f"note: {format_still_structuring_note(names)}\n{text_out}" return ToolResult( - text=rendered.text, + text=text_out, payload={ "rows": visible_rows, "details": { diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py index fad4f97c2..ed55ce9fc 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py @@ -12,6 +12,16 @@ from sqlalchemy import select from shared.models.database.job_result import JobResult +from shared.services.jobs.lifecycle.scope_enrichments import ( + promote_scopes_on_read, +) +from shared.services.retrieval.agent_tools.enrichment_state import ( + STATUS_TEXT, + format_assets_not_refined_note, + format_still_structuring_note, + inspect_scope_page_states, + load_revision_enrichment_status, +) from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, @@ -106,6 +116,33 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ) ).first() if row is None: + revision_id = str( + CorpusRevisionContext.resolve_revision( + doc.document_id, doc.current_job_result_id + ) + or "" + ) + if revision_id: + rev_status = await load_revision_enrichment_status(ctx.db, [revision_id]) + if rev_status.get(revision_id) == STATUS_TEXT: + note = format_still_structuring_note( + [doc.source_file_name or document_id] + ) + return ToolResult(text=note, error=note) + summary = await inspect_scope_page_states( + ctx.db, + namespace=ctx.namespace, + document_id=document_id, + job_result_id=revision_id, + ) + if summary.has_unrefined_pages: + if summary.page_rows: + await promote_scopes_on_read(ctx.db, list(summary.page_rows)) + note = format_assets_not_refined_note( + summary.page_range_label, + state="text" if "text" in summary.states else "coarse", + ) + return ToolResult(text=note, error=note) return ToolResult(text="", error=f"unknown table chunk_id: {chunk_id} in {document_id}") chunk, section_path = row try: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py index d36dd841e..13da12933 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py @@ -70,6 +70,9 @@ section_path_received, section_path_subtree_filter, ) +from shared.services.jobs.lifecycle.scope_enrichments import ( + promote_scopes_on_read, +) from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path _SAME_AS_MARKER_RE = re.compile(r"\[SAME-AS (.+?) p(\d+)\]") @@ -221,6 +224,27 @@ async def _resolve_same_as_markers( row["content"] = content +def _page_image_display_content(row: dict[str, Any]) -> str: + metadata = row.get("chunk_metadata") or row.get("metadata") or {} + raw_assets = metadata.get("page_assets") if isinstance(metadata, dict) else None + lines: list[str] = [] + if isinstance(raw_assets, list): + for item in raw_assets: + if not isinstance(item, dict): + continue + ref = ( + str(item.get("asset_url") or "").strip() + or str(item.get("artifact_ref") or "").strip() + ) + if ref: + lines.append(f"[Page image: {ref}]") + if not lines: + fallback_ref = str(row.get("asset_url") or "").strip() + if fallback_ref: + lines.append(f"[Page image: {fallback_ref}]") + return "\n".join(lines) + + def _compose_explore_text( row: dict[str, Any], rows_by_chunk_id: dict[str, dict[str, Any]], @@ -228,6 +252,13 @@ def _compose_explore_text( char_budget: int, ) -> str: base_content = str(row.get("content") or "") + if ( + not base_content.strip() + and normalize_chunk_type(row.get("chunk_type")) == "page" + ): + page_display = _page_image_display_content(row) + if page_display: + return page_display display = _explore_display_by_target( row, rows_by_chunk_id, @@ -273,11 +304,31 @@ def _explore_display_by_target( def _https_image_media(row: dict[str, Any]) -> list[dict[str, str]]: - if normalize_chunk_type(row.get("chunk_type")) != "image": + chunk_type = normalize_chunk_type(row.get("chunk_type")) + if chunk_type == "image": + url = str(row.get("asset_url") or "").strip() + if url.startswith("https://"): + return [{"type": "image_url", "url": url}] return [] - url = str(row.get("asset_url") or "").strip() - if url.startswith("https://"): - return [{"type": "image_url", "url": url}] + raw_content = str(row.get("content") or "").strip() + if chunk_type == "page" and ( + not raw_content or raw_content == _page_image_display_content(row) + ): + media: list[dict[str, str]] = [] + metadata = row.get("chunk_metadata") or row.get("metadata") or {} + raw_assets = metadata.get("page_assets") if isinstance(metadata, dict) else None + if isinstance(raw_assets, list): + for item in raw_assets: + if not isinstance(item, dict): + continue + url = str(item.get("asset_url") or "").strip() + if url.startswith("https://"): + media.append({"type": "image_url", "url": url}) + if not media: + url = str(row.get("asset_url") or "").strip() + if url.startswith("https://"): + media.append({"type": "image_url", "url": url}) + return media return [] @@ -724,7 +775,14 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: char_budget=char_budget, ) if include_assets - else row.get("content") + else ( + row.get("content") + or ( + _page_image_display_content(row) + if chunk_type == "page" + else "" + ) + ) ) elif chunk_type == "table": composed["content"] = render_explore_table( @@ -734,6 +792,8 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: composed["content"] = _image_display_content(row) assembled.append(composed) + await promote_scopes_on_read(ctx.db, base_rows) + lines = ["refs:", *(f" {line}" for line in status_lines)] for row in assembled: lines.append( diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py index d265f163b..3ac70fe30 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py @@ -1,8 +1,9 @@ """``corpus.recall`` — fuzzy ranked candidate search. -One live scoring path: ``search.map_unit_discovery.map_unit_discovery`` -(persisted map-unit BM25 over path + content, already RRF-fused there). -Exact-string lookup is ``corpus.grep``, not a second channel here. +One live scoring path: ``search.visual_discovery.hybrid_discovery`` — +persisted map-unit BM25 over path + content (already RRF-fused in +``search.map_unit_discovery``) plus visual page-embedding hits for scanned +pages. Exact-string lookup is ``corpus.grep``, not a second channel here. ``scope`` is applied in that same discovery SQL: a named section is that section and everything under it, scored before ranking. Classic retrieval @@ -37,7 +38,7 @@ from shared.services.retrieval.agent_tools.snippet import build_row, build_snippet, format_row from shared.services.retrieval.hydration.row_utils import normalize_chunk_type from shared.services.retrieval.settings import ASSET_CHUNK_TYPES -from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery +from shared.services.retrieval.search.visual_discovery import hybrid_discovery _DEFAULT_LIMIT = 10 @@ -47,13 +48,21 @@ def _identifier_snippet(row: dict[str, Any]) -> str: snippet = str(row.get("snippet") or "").strip() if snippet: return snippet - if normalize_chunk_type(row.get("chunk_type")) == "table": + chunk_type = normalize_chunk_type(row.get("chunk_type")) + if chunk_type == "table": metadata = row.get("chunk_metadata") or {} if not isinstance(metadata, dict): metadata = {} summary = str(metadata.get("summary") or "").strip() return build_snippet(summary) if summary else "" - return build_snippet(str(row.get("content") or "")) + text_snippet = build_snippet(str(row.get("content") or "")) + if text_snippet: + return text_snippet + if chunk_type == "page": + metadata = row.get("chunk_metadata") or {} + if isinstance(metadata, dict) and metadata.get("page_assets"): + return "(scanned page image — call corpus.read to inspect page image)" + return "" @register_tool( @@ -116,7 +125,7 @@ async def recall(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: document_scope = ( ctx.document_scope.narrow(scope_document_ids(scope)) if scope else ctx.document_scope ) - discovery = await map_unit_discovery( + discovery = await hybrid_discovery( ctx.db, user_id=ctx.user_id, namespace=ctx.namespace, diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index c26bdc181..bc4669818 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -153,6 +153,27 @@ async def bump_retrieval_namespace_cache_version( ) +def bump_retrieval_namespace_cache_version_sync( + *, user_id: str, namespace: str +) -> int | None: + try: + from shared.services.redis.redis_sync_service import ( + SyncRedisServiceFactory, + ) + + redis_service = SyncRedisServiceFactory.get_service() + return int( + redis_service.incr( + _namespace_version_key(user_id=user_id, namespace=namespace) + ) + ) + except Exception as exc: + logger.warning( + f"Failed to invalidate retrieval cache namespace (ignored): user_id={user_id}, namespace={namespace}, error={exc}" + ) + return None + + async def invalidate_retrieval_cache_namespaces( *, user_id: str, namespaces: list[str] ) -> None: diff --git a/packages/shared-python/shared/services/retrieval/corpus_storage.py b/packages/shared-python/shared/services/retrieval/corpus_storage.py index e411b6e11..aadfd685b 100644 --- a/packages/shared-python/shared/services/retrieval/corpus_storage.py +++ b/packages/shared-python/shared/services/retrieval/corpus_storage.py @@ -77,6 +77,16 @@ def DocumentMapUnitIndex( else private.DocumentMapUnitIndex ) + @property + def DocumentPageEmbedding( + self, + ) -> type[private.DocumentPageEmbedding] | type[demo.DemoDocumentPageEmbedding]: + return ( + demo.DemoDocumentPageEmbedding + if self.is_demo + else private.DocumentPageEmbedding + ) + @property def RetrievalNamespaceGeneration( self, @@ -149,6 +159,7 @@ def resolve_document(cls, document_id: str) -> CorpusStorage: "retrieval_namespace_generations", "document_map_unit_indexes", "document_map_unit_tokens", + "document_page_embeddings", "retrieval_hit_stats", "document_map_units", "document_sections", diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index e0cd1a5ae..6c8937d0c 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -22,12 +22,12 @@ from shared.services.retrieval.routing import config as query_router_config from shared.services.retrieval.routing.classifier import classify_query_mode from shared.services.retrieval.routing.types import QueryMode, RouterDecision -from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( count_scoped_chunks, load_all_scoped_chunks, ) +from shared.services.retrieval.search.visual_discovery import hybrid_discovery def open_fresh_database_context() -> AbstractAsyncContextManager[AsyncSession]: @@ -182,7 +182,7 @@ async def _try_run_small_corpus_route( async def _run_classic_topk_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: - discovery_result = await map_unit_discovery( + discovery_result = await hybrid_discovery( context.db, user_id=context.user_id, namespace=context.namespace, diff --git a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py index d0f0bd713..8764f8608 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py +++ b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py @@ -22,6 +22,7 @@ from shared.services.retrieval.hydration.row_utils import ( extract_page_nums, normalize_chunk_type, + page_enrichment_state, page_summary, ) from shared.services.retrieval.hydration.table_grid import ( @@ -131,6 +132,11 @@ def _indent_member_part(part: dict[str, Any]) -> dict[str, Any]: def _compose_page_parts(row: dict[str, Any]) -> list[dict[str, Any]]: parts: list[dict[str, Any]] = [] + if page_enrichment_state(row) in {"text", "coarse", "queued"}: + raw_text = str(row.get("content") or "").strip() + if raw_text: + parts.append(_text_part(raw_text)) + return parts summary = page_summary(row) if summary: parts.append(_text_part(summary)) diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index f62f6745c..0edca85c8 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -14,6 +14,7 @@ filter_excluded_rows, iter_connected_target_ids, normalize_chunk_type, + page_enrichment_state, page_summary, ) @@ -68,8 +69,14 @@ async def assemble_retrieval_results( base_content = str(row.get('content') or '') chunk_type = normalize_chunk_type(row.get('chunk_type')) if chunk_type == 'page': - assembled_row['content'] = page_summary(row) - assembled_row['content_source'] = 'summary' + if page_enrichment_state(row) in {'text', 'coarse', 'queued'}: + assembled_row['content'] = base_content + assembled_row['content_source'] = ( + 'raw' if base_content.strip() else 'page_image' + ) + else: + assembled_row['content'] = page_summary(row) + assembled_row['content_source'] = 'summary' page_nums = extract_page_nums(row) if page_nums is not None: assembled_row['page_nums'] = page_nums diff --git a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py index b8f6b8d4c..1eb8b3726 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py +++ b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py @@ -53,6 +53,13 @@ def page_summary(row: dict[str, Any]) -> str: return str(metadata.get('summary') or '').strip() +def page_enrichment_state(row: dict[str, Any]) -> str: + metadata = row.get('chunk_metadata') or row.get('metadata') or {} + if not isinstance(metadata, dict): + return '' + return str(metadata.get('enrichment_state') or '').strip().lower() + + def build_reference_lookup_key( *, document_id: object, diff --git a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py b/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py deleted file mode 100644 index 162b85c25..000000000 --- a/packages/shared-python/shared/services/retrieval/namespace_map_snapshot_redis.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Redis cache for compressed namespace MAP routing snapshots.""" - -from __future__ import annotations - -from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace -from shared.services.redis import RedisServiceFactory - -_CACHE_TTL_SECONDS = 3600 -_KEY_PREFIX = "retrieval:snapshot:v2" - - -class NamespaceMapSnapshotRedisCache: - """Access generation-scoped compressed namespace snapshots in Redis.""" - - @staticmethod - def _build_key(*, user_id: str, namespace: str, generation: int) -> str: - normalized_namespace = normalize_retrieval_namespace(namespace) - return f"{_KEY_PREFIX}:{user_id}:{normalized_namespace}:g{int(generation)}" - - @classmethod - async def get( - cls, *, user_id: str, namespace: str, generation: int - ) -> bytes | None: - service = RedisServiceFactory.get_service() - return await service.get_bytes( - cls._build_key( - user_id=user_id, namespace=namespace, generation=generation - ) - ) - - @classmethod - async def set( - cls, *, user_id: str, namespace: str, generation: int, payload_zlib: bytes - ) -> bool: - service = RedisServiceFactory.get_service() - return await service.set_bytes( - cls._build_key( - user_id=user_id, namespace=namespace, generation=generation - ), - payload_zlib, - ex=_CACHE_TTL_SECONDS, - ) diff --git a/packages/shared-python/shared/services/retrieval/page_embedding_index.py b/packages/shared-python/shared/services/retrieval/page_embedding_index.py new file mode 100644 index 000000000..63e2b5fc8 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/page_embedding_index.py @@ -0,0 +1,205 @@ +"""Publication-time persistence and same-revision rebinding of page embeddings.""" + +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.core.config import settings +from shared.services.ai.embedding import resolve_embedding_model +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.publication_models import DocumentPublicationScope + +__all__ = [ + "extract_and_strip_page_embeddings", + "inherit_missing_page_assets", + "load_revision_page_assets_by_page", + "replace_document_page_embeddings", +] + + +def _coerce_page_nums(raw_page_nums: object) -> list[int]: + if isinstance(raw_page_nums, list): + items = raw_page_nums + elif raw_page_nums is None: + return [] + else: + items = str(raw_page_nums).split(",") + pages: list[int] = [] + for item in items: + try: + page = int(str(item).strip()) + except (TypeError, ValueError): + continue + if page > 0: + pages.append(page) + return pages + + +def extract_and_strip_page_embeddings( + chunk_metadata: dict[str, Any], +) -> dict[int, list[float]]: + """Pop transient ``_page_embedding`` from ``chunk_metadata`` and key by page_num.""" + raw_embedding = chunk_metadata.pop("_page_embedding", None) + if not isinstance(raw_embedding, list) or not raw_embedding: + return {} + try: + vector = [float(val) for val in raw_embedding] + except (TypeError, ValueError): + return {} + pages = _coerce_page_nums(chunk_metadata.get("page_nums")) + if not pages: + return {} + return {pages[0]: vector} + + +def load_revision_page_assets_by_page( + db: Session, + *, + scope: DocumentPublicationScope, +) -> dict[int, dict[str, Any]]: + """Load existing ``page_assets`` keyed by ``page_num`` before chunk replacement.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) + metadatas = db.scalars( + select(corpusStorage.DocumentChunk.chunk_metadata) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) + .where(corpusStorage.DocumentChunk.chunk_type == "page") + .order_by( + corpusStorage.DocumentChunk.sort_order, + corpusStorage.DocumentChunk.id, + ) + ).all() + assets_by_page: dict[int, dict[str, Any]] = {} + for metadata in metadatas: + if not isinstance(metadata, dict): + continue + raw_assets = metadata.get("page_assets") + if not isinstance(raw_assets, list): + continue + for item in raw_assets: + if not isinstance(item, dict): + continue + try: + page_num = int(item.get("page_num")) # type: ignore[arg-type] + except (TypeError, ValueError): + continue + if page_num > 0 and page_num not in assets_by_page: + assets_by_page[page_num] = dict(item) + return assets_by_page + + +def inherit_missing_page_assets( + chunk_metadata: dict[str, Any], + page_assets_by_page: dict[int, dict[str, Any]], +) -> None: + """Attach prior M0/coarse ``page_assets`` when a restructured page chunk lacks them.""" + if not page_assets_by_page: + return + existing = chunk_metadata.get("page_assets") + if isinstance(existing, list) and existing: + return + pages = _coerce_page_nums(chunk_metadata.get("page_nums")) + inherited = [ + dict(page_assets_by_page[page]) + for page in pages + if page in page_assets_by_page + ] + if inherited: + chunk_metadata["page_assets"] = inherited + + +def _build_owner_chunk_by_page( + db: Session, + *, + scope: DocumentPublicationScope, +) -> dict[int, str]: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) + rows = db.execute( + select( + corpusStorage.DocumentChunk.chunk_id, + corpusStorage.DocumentChunk.chunk_metadata, + ) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) + .where(corpusStorage.DocumentChunk.chunk_type == "page") + .order_by( + corpusStorage.DocumentChunk.sort_order, + corpusStorage.DocumentChunk.id, + ) + ).all() + owner_by_page: dict[int, str] = {} + for chunk_id, metadata in rows: + if not chunk_id or not isinstance(metadata, dict): + continue + for page_num in _coerce_page_nums(metadata.get("page_nums")): + owner_by_page.setdefault(page_num, str(chunk_id)) + return owner_by_page + + +def replace_document_page_embeddings( + db: Session, + *, + scope: DocumentPublicationScope, + incoming_embeddings_by_page: dict[int, list[float]] | None = None, +) -> int: + """Persist new M0 page embeddings, or rebind existing ones to the current revision's page chunks.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) + owner_by_page = _build_owner_chunk_by_page(db, scope=scope) + if not owner_by_page and not incoming_embeddings_by_page: + return 0 + + if incoming_embeddings_by_page: + model_name = resolve_embedding_model(getattr(settings, "EMBEDDING_MODEL", None)) + db.execute( + delete(corpusStorage.DocumentPageEmbedding) + .where(corpusStorage.DocumentPageEmbedding.document_id == scope.document_id) + .where( + corpusStorage.DocumentPageEmbedding.job_result_id == scope.job_result_id + ) + ) + inserted = 0 + for page_num in sorted(incoming_embeddings_by_page): + vector = incoming_embeddings_by_page[page_num] + chunk_id = owner_by_page.get(page_num) + if not chunk_id or not vector: + continue + db.add( + corpusStorage.DocumentPageEmbedding( + id=f"dpe_{uuid4().hex[:12]}", + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + chunk_id=chunk_id, + page_num=int(page_num), + modality="visual", + model=model_name, + dims=len(vector), + embedding=vector, + ) + ) + inserted += 1 + db.flush() + return inserted + + current_rows = list( + db.scalars( + select(corpusStorage.DocumentPageEmbedding) + .where(corpusStorage.DocumentPageEmbedding.document_id == scope.document_id) + .where( + corpusStorage.DocumentPageEmbedding.job_result_id == scope.job_result_id + ) + .order_by(corpusStorage.DocumentPageEmbedding.page_num) + ).all() + ) + for row in current_rows: + target_chunk_id = owner_by_page.get(int(row.page_num)) + if target_chunk_id and row.chunk_id != target_chunk_id: + row.chunk_id = target_chunk_id + if current_rows: + db.flush() + return len(current_rows) diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 13d60f17f..32b5f37db 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, cast from uuid import uuid4 -from sqlalchemy import delete +from sqlalchemy import delete, select from sqlalchemy.orm import Session from shared.models.database.document import ( @@ -13,6 +13,10 @@ DocumentSection, ) from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.page_embedding_index import ( + extract_and_strip_page_embeddings, + replace_document_page_embeddings, +) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.publication_trace_stage import trace_publication_stage from shared.services.retrieval.serving_manifest import persist_revision_serving_state @@ -72,14 +76,19 @@ def replace_document_revision_content( db=db, scope=scope, section_summaries=section_summaries, + preload_existing=False, ) prepared_chunks: list[ tuple[int, dict[str, Any], dict[str, Any], str | None, DocumentSection] ] = [] + incoming_embeddings_by_page: dict[int, list[float]] = {} with trace_publication_stage(trace, "sections_prepare"): for index, chunk in enumerate(chunks): safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) - chunk_metadata = _get_chunk_metadata(safe_chunk) + chunk_metadata = dict(_get_chunk_metadata(safe_chunk)) + incoming_embeddings_by_page.update( + extract_and_strip_page_embeddings(chunk_metadata) + ) source_path = _get_source_path( chunk=safe_chunk, chunk_metadata=chunk_metadata, @@ -124,6 +133,11 @@ def replace_document_revision_content( trace=trace, clear_existing=False, ) + replace_document_page_embeddings( + db, + scope=scope, + incoming_embeddings_by_page=incoming_embeddings_by_page or None, + ) manifest_payload = persist_revision_serving_state(db, scope=scope, trace=trace) return manifest_payload @@ -135,6 +149,7 @@ def __init__( db: Session, scope: DocumentPublicationScope, section_summaries: dict[str, str] | None = None, + preload_existing: bool = True, ) -> None: self._db = db self._scope = scope @@ -143,17 +158,61 @@ def __init__( remove_nul_characters(section_summaries or {}), ) self._sections_by_path: dict[str, DocumentSection] = {} + self._preexisting_paths: set[str] = set() + self._existing_loaded = not preload_existing + + def _ensure_existing_loaded(self) -> None: + if self._existing_loaded: + return + self._existing_loaded = True + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace( + self._scope.namespace + ) + existing_sections = self._db.scalars( + select(corpusStorage.DocumentSection) + .where( + corpusStorage.DocumentSection.document_id + == self._scope.document_id + ) + .where( + corpusStorage.DocumentSection.job_result_id + == self._scope.job_result_id + ) + .order_by( + corpusStorage.DocumentSection.sort_order, + corpusStorage.DocumentSection.section_id, + ) + ).all() + for existing in existing_sections: + self._sections_by_path.setdefault(existing.section_path, existing) + self._preexisting_paths.add(existing.section_path) def ensure_section(self, section_path: str) -> DocumentSection: + self._ensure_existing_loaded() corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(self._scope.namespace) existing_section = self._sections_by_path.get(section_path) if existing_section is not None: + if ( + summary := self._section_summaries.get(section_path) + ) and ( + section_path not in self._preexisting_paths + or not existing_section.summary + ): + existing_section.summary = summary return existing_section path_parts = [part for part in section_path.split(" / ") if part] for depth in range(1, len(path_parts) + 1): ancestor_path = " / ".join(path_parts[:depth]) - if ancestor_path in self._sections_by_path: + existing_ancestor = self._sections_by_path.get(ancestor_path) + if existing_ancestor is not None: + if ( + summary := self._section_summaries.get(ancestor_path) + ) and ( + ancestor_path not in self._preexisting_paths + or not existing_ancestor.summary + ): + existing_ancestor.summary = summary continue ancestor_section = corpusStorage.DocumentSection( diff --git a/packages/shared-python/shared/services/retrieval/publication_scope_content.py b/packages/shared-python/shared/services/retrieval/publication_scope_content.py new file mode 100644 index 000000000..b3236fbc2 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_scope_content.py @@ -0,0 +1,310 @@ +"""Scope-granular retrieval content replacement for Phase B page-memory enrichment.""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Any, cast + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.cache_service import ( + bump_retrieval_namespace_cache_version_sync, +) +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.map_unit_index import replace_document_map_units +from shared.services.retrieval.page_embedding_index import ( + extract_and_strip_page_embeddings, + inherit_missing_page_assets, + load_revision_page_assets_by_page, + replace_document_page_embeddings, +) +from shared.services.retrieval.publication_chunk_copy import insert_chunks_with_copy +from shared.services.retrieval.publication_content import ( + DocumentSectionPublisher, + _build_document_chunk, + _get_chunk_metadata, + _get_source_path, + deduplicate_chunks_by_source_path, +) +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.publication_trace_stage import trace_publication_stage +from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path +from shared.services.retrieval.serving_manifest import persist_revision_serving_state +from shared.utils.json_utils import remove_nul_characters + +if TYPE_CHECKING: + from shared.services.jobs.lifecycle.publication_trace import PublicationTrace + +_SCOPE_ID_RE = re.compile(r"^p([1-9]\d*)-([1-9]\d*)$") +_REFINED_ENRICHMENT_STATE = "refined" + + +def replace_document_scope_content( + db: Session, + *, + scope: DocumentPublicationScope, + scope_id: str, + chunks: list[dict[str, Any]], + section_summaries: dict[str, str] | None = None, + trace: PublicationTrace | None = None, +) -> dict[str, Any]: + """Replace one scope's sections/chunks, rebuild map units, and bump cache version.""" + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) + with trace_publication_stage(trace, "chunks_prepare"): + deduped_chunks = deduplicate_chunks_by_source_path(chunks) + prior_page_assets_by_page = load_revision_page_assets_by_page( + db, + scope=scope, + ) + + with trace_publication_stage(trace, "chunks_persist"): + delete_stmt = ( + delete(corpusStorage.DocumentChunk) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) + ) + if scope_id != "*": + delete_stmt = delete_stmt.where( + corpusStorage.DocumentChunk.chunk_metadata["scope_id"].as_string() + == scope_id + ) + db.execute(delete_stmt) + + with trace_publication_stage(trace, "sections_persist"): + _delete_scope_exclusive_leaf_sections( + db, + scope=scope, + incoming_chunks=deduped_chunks, + ) + + section_publisher = DocumentSectionPublisher( + db=db, + scope=scope, + section_summaries=section_summaries, + preload_existing=True, + ) + prepared_chunks: list[ + tuple[int, dict[str, Any], dict[str, Any], str | None, DocumentSection] + ] = [] + incoming_embeddings_by_page: dict[int, list[float]] = {} + with trace_publication_stage(trace, "sections_prepare"): + for index, chunk in enumerate(deduped_chunks): + safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) + chunk_metadata = dict(_get_chunk_metadata(safe_chunk)) + incoming_embeddings_by_page.update( + extract_and_strip_page_embeddings(chunk_metadata) + ) + if scope_id != "*": + chunk_metadata["scope_id"] = scope_id + chunk_type = str( + safe_chunk.get("type") or safe_chunk.get("chunk_type") or "text" + ) + if chunk_type == "page": + if scope_id != "*": + chunk_metadata["enrichment_state"] = _REFINED_ENRICHMENT_STATE + else: + chunk_metadata.setdefault( + "enrichment_state", _REFINED_ENRICHMENT_STATE + ) + inherit_missing_page_assets( + chunk_metadata, + prior_page_assets_by_page, + ) + source_path = _get_source_path( + chunk=safe_chunk, + chunk_metadata=chunk_metadata, + ) + section_path = section_path_from_chunk_path( + source_path, + source_file_name=scope.source_file_name, + ) + section = section_publisher.ensure_section(section_path) + prepared_chunks.append( + (index, safe_chunk, chunk_metadata, source_path, section) + ) + with trace_publication_stage(trace, "sections_persist"): + db.flush() + + chunk_rows: list[DocumentChunk] = [] + with trace_publication_stage(trace, "chunks_prepare"): + for index, safe_chunk, chunk_metadata, source_path, section in prepared_chunks: + chunk_rows.append( + _build_document_chunk( + chunk=safe_chunk, + chunk_metadata=chunk_metadata, + source_path=source_path, + section=section, + scope=scope, + fallback_sort_order=index, + ) + ) + with trace_publication_stage(trace, "chunks_persist"): + if chunk_rows: + insert_chunks_with_copy(db, chunk_rows) + _normalize_revision_sort_orders(db, scope=scope) + + replace_document_map_units( + db, + scope=scope, + trace=trace, + clear_existing=True, + ) + replace_document_page_embeddings( + db, + scope=scope, + incoming_embeddings_by_page=incoming_embeddings_by_page or None, + ) + manifest_payload = persist_revision_serving_state(db, scope=scope, trace=trace) + bump_retrieval_namespace_cache_version_sync( + user_id=scope.user_id, + namespace=scope.namespace, + ) + return manifest_payload + + +def _delete_scope_exclusive_leaf_sections( + db: Session, + *, + scope: DocumentPublicationScope, + incoming_chunks: list[dict[str, Any]], +) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) + other_section_ids = set( + db.scalars( + select(corpusStorage.DocumentChunk.section_id) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) + .where(corpusStorage.DocumentChunk.section_id.is_not(None)) + .distinct() + ).all() + ) + existing_sections = list( + db.scalars( + select(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == scope.document_id) + .where(corpusStorage.DocumentSection.job_result_id == scope.job_result_id) + ).all() + ) + if not existing_sections: + return + + retained_paths: set[str] = set() + for section in existing_sections: + if section.section_id in other_section_ids: + retained_paths.update(_self_and_ancestor_paths(section.section_path)) + + for chunk in incoming_chunks: + safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) + chunk_metadata = _get_chunk_metadata(safe_chunk) + source_path = _get_source_path( + chunk=safe_chunk, + chunk_metadata=chunk_metadata, + ) + section_path = section_path_from_chunk_path( + source_path, + source_file_name=scope.source_file_name, + ) + retained_paths.update(_strict_ancestor_paths(section_path)) + + stale_section_ids = [ + section.section_id + for section in existing_sections + if section.section_path not in retained_paths + ] + if stale_section_ids: + db.execute( + delete(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == scope.document_id) + .where(corpusStorage.DocumentSection.job_result_id == scope.job_result_id) + .where(corpusStorage.DocumentSection.section_id.in_(stale_section_ids)) + ) + + +def _normalize_revision_sort_orders( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(scope.namespace) + chunks = list( + db.scalars( + select(corpusStorage.DocumentChunk) + .where(corpusStorage.DocumentChunk.document_id == scope.document_id) + .where(corpusStorage.DocumentChunk.job_result_id == scope.job_result_id) + .order_by( + corpusStorage.DocumentChunk.sort_order, + corpusStorage.DocumentChunk.id, + ) + ).all() + ) + if not chunks: + return + + ordered_chunks = sorted( + chunks, + key=lambda row: ( + _scope_sort_key(row.chunk_metadata), + int(row.sort_order or 0), + str(row.id or ""), + ), + ) + for new_order, chunk in enumerate(ordered_chunks): + if chunk.sort_order != new_order: + chunk.sort_order = new_order + + sections = list( + db.scalars( + select(corpusStorage.DocumentSection) + .where(corpusStorage.DocumentSection.document_id == scope.document_id) + .where(corpusStorage.DocumentSection.job_result_id == scope.job_result_id) + .order_by( + corpusStorage.DocumentSection.sort_order, + corpusStorage.DocumentSection.section_id, + ) + ).all() + ) + sections_by_id = {section.section_id: section for section in sections} + sections_by_path = {section.section_path: section for section in sections} + visited_paths: list[str] = [] + seen_paths: set[str] = set() + for chunk in ordered_chunks: + section = sections_by_id.get(chunk.section_id or "") + if section is None: + continue + for ancestor_path in _self_and_ancestor_paths(section.section_path): + if ancestor_path not in seen_paths and ancestor_path in sections_by_path: + seen_paths.add(ancestor_path) + visited_paths.append(ancestor_path) + for section in sections: + if section.section_path not in seen_paths: + seen_paths.add(section.section_path) + visited_paths.append(section.section_path) + for new_order, path in enumerate(visited_paths): + section = sections_by_path[path] + if section.sort_order != new_order: + section.sort_order = new_order + + db.flush() + + +def _scope_sort_key(metadata: Any) -> tuple[int, int, int]: + if isinstance(metadata, dict): + scope_id = str(metadata.get("scope_id") or "").strip() + matched = _SCOPE_ID_RE.fullmatch(scope_id) + if matched is not None: + return (0, int(matched.group(1)), int(matched.group(2))) + return (1, 0, 0) + + +def _self_and_ancestor_paths(section_path: str) -> list[str]: + parts = [part for part in section_path.split(" / ") if part] + return [" / ".join(parts[:depth]) for depth in range(1, len(parts) + 1)] + + +def _strict_ancestor_paths(section_path: str) -> list[str]: + parts = [part for part in section_path.split(" / ") if part] + if len(parts) <= 1: + return [] + return [" / ".join(parts[:depth]) for depth in range(1, len(parts))] diff --git a/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py index 7e864e831..8e18a4882 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py @@ -18,6 +18,7 @@ RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 CHANNEL_WEIGHT_CONTENT = 2.0 +CHANNEL_WEIGHT_VISUAL = 1.5 # Persisted map-unit token format. Bump when tokenizer semantics change. # v1: character-level CJK regex. v2: word-level jieba (text_utils). MAP_UNIT_INDEX_FORMAT_VERSION = 2 diff --git a/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py b/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py index 34c3b5c6f..6659e26aa 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py +++ b/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py @@ -1,7 +1,6 @@ """Publication-time hierarchy provider over section/chunk rows. -Extracted from the map-nav package so document publication and classic -recall can build map-unit indexes without the episode kernel. +Document publication and classic recall use this to build map-unit indexes. """ from __future__ import annotations diff --git a/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py b/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py index 0ad4b85cb..00df4db8d 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py +++ b/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py @@ -1,14 +1,7 @@ """Score-based folding for an oversized ``corpus.outline``/``corpus.node_filter`` map. -Adapted from the archived MAP-NAV pooling/hiding algorithm -(``deprecated/mapnav/nav/nav_map_scores.py``'s ``_pool_unit_scores_to_tree`` -and ``nav_projection.py``'s ``_apply_budget_hide``) onto the live, -``document_sections``-based tree ``outline``/``node_filter`` already build — -not onto the retired ``ToolSpace``/``LazyKnowhereProvider`` abstraction those -files used. The deprecated files are unmodified; this is a fresh module using -the same idea against live data. - -Only invoked when a rendered map exceeds +Folds the ``document_sections`` tree that ``outline``/``node_filter`` already +build. Only invoked when a rendered map exceeds ``agent_tools.registry.MAP_TOOL_CHAR_BUDGET``. Scoring needs the end user's original query (``ToolContext.query``) — see that field's docstring for why a caller without one gets a "narrow the scope" error instead of a silently @@ -298,9 +291,8 @@ def fold_map_nodes( Every id in ``keep_ids`` and all of its ancestors are never hidden; a subtree is a removal candidate only when no kept row lies in it. - Candidates go lowest score first; ties follow the archived MAP-NAV - order (``_apply_budget_hide``): most descendants first, then deepest - first, then ``section_id``. The budget counts the placeholder rows left + Candidates go lowest score first; ties go most descendants first, then + deepest first, then ``section_id``. The budget counts the placeholder rows left behind, not just the visible rows. ``overflow`` is set when the rows that must stay (plus their placeholders) still exceed ``char_budget``; the caller fails the call rather than hiding a kept row or cutting text. diff --git a/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py b/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py index 91bfa047b..c2c3e383d 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py +++ b/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py @@ -1,7 +1,6 @@ """Predicate compile/match for section path/summary filters. -Extracted from the map-nav tree walker. Live consumers evaluate these -predicates against persisted section rows, not an in-memory episode tree. +Live consumers evaluate these predicates against persisted section rows. """ from __future__ import annotations diff --git a/packages/shared-python/shared/services/retrieval/search/visual_discovery.py b/packages/shared-python/shared/services/retrieval/search/visual_discovery.py new file mode 100644 index 000000000..3415f2a2f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/visual_discovery.py @@ -0,0 +1,460 @@ +"""Visual page embedding discovery (pgvector cosine) and RRF hybrid fusion.""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Mapping +from typing import Any + +from loguru import logger +from pgvector import HalfVector +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.config import settings +from shared.services.ai.embedding import embed_query +from shared.services.retrieval.corpus_storage import CorpusStorage +from shared.services.retrieval.document_scope import DocumentScope +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.row_utils import ( + iter_connected_target_ids, + normalize_chunk_type, +) +from shared.services.retrieval.scoring.knowhere_hybrid import ( + CHANNEL_WEIGHT_CONTENT, + CHANNEL_WEIGHT_PATH, + CHANNEL_WEIGHT_VISUAL, + RRF_K, +) +from shared.services.retrieval.search.map_unit_discovery import ( + DiscoveryResult, + _as_metadata_dict, + _build_signal_clause, + map_unit_discovery, +) +from shared.services.retrieval.search.scoring import normalize_row_scores +from shared.services.retrieval.settings import ASSET_CHUNK_TYPES + +__all__ = ["hybrid_discovery", "visual_page_discovery"] + + +def _build_visual_revision_scope( + revision_pins: Mapping[str, str] | None, +) -> tuple[str, str, dict[str, Any]]: + if revision_pins is None: + return ("AND d.current_job_result_id = dpe.job_result_id", "", {}) + pairs = [ + (str(document_id).strip(), str(job_result_id).strip()) + for document_id, job_result_id in revision_pins.items() + if str(document_id).strip() and str(job_result_id).strip() + ] + if not pairs: + return "", "AND FALSE", {} + params: dict[str, Any] = {} + placeholders: list[str] = [] + for index, (document_id, job_result_id) in enumerate(pairs): + document_key = f"_vpin_document_{index}" + revision_key = f"_vpin_revision_{index}" + placeholders.append(f"(:{document_key}, :{revision_key})") + params[document_key] = document_id + params[revision_key] = job_result_id + return ( + "", + f"AND (dpe.document_id, dpe.job_result_id) IN ({', '.join(placeholders)})", + params, + ) + + +def _build_visual_section_subtree_clause( + targets: list[tuple[str, str | None]] | None, +) -> tuple[str, dict[str, Any]]: + if not targets or not any(section_path for _document_id, section_path in targets): + return "", {} + parts: list[str] = [] + params: dict[str, Any] = {} + for index, (document_id, section_path) in enumerate(targets): + doc_key = f"_vsec_doc_{index}" + params[doc_key] = document_id + if not section_path: + parts.append(f"(dpe.document_id = :{doc_key})") + continue + path_key = f"_vsec_path_{index}" + like_key = f"_vsec_pathlike_{index}" + params[path_key] = section_path + params[like_key] = f"{section_path} / %" + parts.append( + f"(dpe.document_id = :{doc_key} AND " + f"(ds.section_path = :{path_key} OR ds.section_path LIKE :{like_key}))" + ) + return f"AND ({' OR '.join(parts)})", params + + +def _build_visual_exclude_sections_clause( + exclude_sections: list[dict[str, str]], +) -> tuple[str, dict[str, Any]]: + parts: list[str] = [] + params: dict[str, Any] = {} + for index, item in enumerate(exclude_sections): + if not isinstance(item, dict): + continue + document_id = str(item.get("document_id") or "").strip() + section_path = str(item.get("section_path") or "").strip() + if not document_id or not section_path: + continue + doc_key = f"_vexc_doc_{index}" + path_key = f"_vexc_path_{index}" + like_key = f"_vexc_pathlike_{index}" + params[doc_key] = document_id + params[path_key] = section_path + params[like_key] = f"{section_path} / %" + parts.append( + f"(dpe.document_id = :{doc_key} AND " + f"(ds.section_path = :{path_key} OR ds.section_path LIKE :{like_key}))" + ) + if not parts: + return "", {} + return f"AND NOT ({' OR '.join(parts)})", params + + +async def visual_page_discovery( + db: AsyncSession | None, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str] | None = None, + exclude_sections: list[dict[str, str]] | None = None, + document_scope: DocumentScope = DocumentScope(), + chunk_types: set[str] | None = None, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, + section_targets: list[tuple[str, str | None]] | None = None, +) -> DiscoveryResult: + """Rank persisted page embeddings via pgvector cosine similarity.""" + if db is None: + raise ValueError("database session required for visual_page_discovery") + if not bool(getattr(settings, "VISUAL_RETRIEVAL_ENABLED", True)): + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + normalized_query = str(query or "").strip() + if not normalized_query or top_k <= 0: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + if revision_pins is not None and not revision_pins: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + if ( + chunk_types is not None + and "page" not in chunk_types + and not chunk_types.issubset(ASSET_CHUNK_TYPES) + ): + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + t0 = time.monotonic() + corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace) + exclude_doc_list = list(exclude_document_ids or []) + exclude_sec_list = list(exclude_sections or []) + + revision_join, revision_clause, revision_params = _build_visual_revision_scope( + revision_pins + ) + scoped_docs = document_scope.excluding(exclude_doc_list) + document_scope_clause, document_scope_params = scoped_docs.sql() + params: dict[str, Any] = { + "user_id": corpusStorage.resolve_owner(user_id), + "namespace": namespace, + **revision_params, + **document_scope_params, + } + + exists_sql = f""" + SELECT EXISTS ( + SELECT 1 + FROM document_page_embeddings dpe + JOIN documents d + ON d.document_id = dpe.document_id + {revision_join} + WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + {revision_clause} + {document_scope_clause} + ) + """ + has_embeddings = bool( + ( + await db.execute(text(corpusStorage.compile_sql(exists_sql)), params) + ).scalar() + ) + if not has_embeddings: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + query_vector = await asyncio.to_thread( + embed_query, + normalized_query, + usage_task="retrieval.query_embed", + ) + vector_literal = HalfVector(query_vector).to_text() + + signal_clause, signal_params = _build_signal_clause( + signal_paths or [], filter_mode + ) + section_clause, section_params = _build_visual_section_subtree_clause( + section_targets + ) + exclude_clause, exclude_params = _build_visual_exclude_sections_clause( + exclude_sec_list + ) + query_params: dict[str, Any] = { + **params, + **signal_params, + **section_params, + **exclude_params, + "query_embedding": vector_literal, + "top_k": int(top_k), + } + + search_sql = f""" + SELECT + ranked.page_num, + ranked.similarity, + ranked.chunk_id, + ranked.document_id, + ranked.section_id, + ranked.chunk_type, + ranked.content, + ranked.source_chunk_path, + ranked.file_path, + ranked.chunk_metadata, + ranked.job_result_id, + ranked.sort_order, + ranked.section_path, + ranked.source_file_name, + ranked.job_id, + ranked.result_raw_prefix + FROM ( + SELECT DISTINCT ON (dpe.document_id, dpe.chunk_id) + dpe.page_num, + 1.0 - (dpe.embedding <=> CAST(:query_embedding AS halfvec)) + AS similarity, + dc.chunk_id, + dc.document_id, + dc.section_id, + dc.chunk_type, + dc.content, + dc.source_chunk_path, + dc.file_path, + dc.chunk_metadata, + dc.job_result_id, + dc.sort_order, + ds.section_path, + d.source_file_name, + jr.job_id, + jr.document_metadata ->> 'result_raw_prefix' AS result_raw_prefix + FROM document_page_embeddings dpe + JOIN documents d + ON d.document_id = dpe.document_id + {revision_join} + JOIN document_chunks dc + ON dc.document_id = dpe.document_id + AND dc.job_result_id = dpe.job_result_id + AND dc.chunk_id = dpe.chunk_id + LEFT JOIN document_sections ds + ON ds.section_id = dc.section_id + LEFT JOIN job_results jr + ON jr.id = dc.job_result_id + WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + {revision_clause} + {document_scope_clause} + {signal_clause} + {section_clause} + {exclude_clause} + ORDER BY dpe.document_id, + dpe.chunk_id, + dpe.embedding <=> CAST(:query_embedding AS halfvec) ASC, + dpe.page_num ASC + ) ranked + ORDER BY ranked.similarity DESC, ranked.page_num ASC + LIMIT :top_k + """ + result = await db.execute( + text(corpusStorage.compile_sql(search_sql)), query_params + ) + + primaries: list[dict[str, Any]] = [] + for raw_row in result.all(): + row = dict(raw_row._mapping) + similarity = max(0.0, float(row.pop("similarity", 0.0) or 0.0)) + row.pop("page_num", None) + row["chunk_metadata"] = _as_metadata_dict(row.get("chunk_metadata")) + row["score"] = similarity + primaries.append(row) + + if chunk_types is not None and chunk_types.issubset(ASSET_CHUNK_TYPES): + connected = await hydrate_connected_target_rows( + db=db, + rows=primaries, + exclude_document_ids=exclude_doc_list, + exclude_sections=exclude_sec_list, + document_scope=scoped_docs, + revision_pins=revision_pins, + ) + connected_by_id = { + str(row.get("chunk_id") or ""): row + for row in connected + if row.get("chunk_id") + } + asset_rows: list[dict[str, Any]] = [] + seen_ids: set[str] = set() + for primary in primaries: + score = float(primary.get("score") or 0.0) + for target_id in iter_connected_target_ids(primary): + target = connected_by_id.get(target_id) + if target is None: + continue + cid = str(target.get("chunk_id") or "") + if ( + not cid + or cid in seen_ids + or normalize_chunk_type(target.get("chunk_type")) not in chunk_types + ): + continue + seen_ids.add(cid) + item = dict(target) + item["score"] = score + asset_rows.append(item) + fused_rows = asset_rows[:top_k] + else: + fused_rows = primaries + + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + latency = int((time.monotonic() - t0) * 1000) + logger.info( + " search.visual_page_discovery: {} fused rows, {}ms", + len(fused_rows), + latency, + ) + return DiscoveryResult( + status="discovery_done", + payload={"fused_rows": fused_rows}, + latency_ms=latency, + ) + + +async def hybrid_discovery( + db: AsyncSession | None, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str] | None = None, + exclude_sections: list[dict[str, str]] | None = None, + document_scope: DocumentScope = DocumentScope(), + chunk_types: set[str] | None = None, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + revision_pins: Mapping[str, str] | None = None, + publish_index_readiness: bool = True, + section_targets: list[tuple[str, str | None]] | None = None, +) -> DiscoveryResult: + """Fuse persisted lexical map-unit BM25 discovery with visual page discovery via RRF.""" + exclude_doc_list = list(exclude_document_ids or []) + exclude_sec_list = list(exclude_sections or []) + + lexical_result = await map_unit_discovery( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_doc_list, + exclude_sections=exclude_sec_list, + document_scope=document_scope, + chunk_types=chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + revision_pins=revision_pins, + publish_index_readiness=publish_index_readiness, + section_targets=section_targets, + ) + if not bool(getattr(settings, "VISUAL_RETRIEVAL_ENABLED", True)): + return lexical_result + + visual_result = await visual_page_discovery( + db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_doc_list, + exclude_sections=exclude_sec_list, + document_scope=document_scope, + chunk_types=chunk_types, + signal_paths=signal_paths, + filter_mode=filter_mode, + revision_pins=revision_pins, + section_targets=section_targets, + ) + + lexical_rows = list(lexical_result.payload.get("fused_rows") or []) + visual_rows = list(visual_result.payload.get("fused_rows") or []) + + if not visual_rows: + return lexical_result + if not lexical_rows: + return visual_result + + lexical_weight = CHANNEL_WEIGHT_PATH + CHANNEL_WEIGHT_CONTENT + rrf_scores: dict[tuple[str, str], float] = {} + rows_by_key: dict[tuple[str, str], dict[str, Any]] = {} + + for rank, row in enumerate(lexical_rows): + key = (str(row.get("document_id") or ""), str(row.get("chunk_id") or "")) + rows_by_key[key] = dict(row) + rrf_scores[key] = rrf_scores.get(key, 0.0) + lexical_weight / ( + RRF_K + rank + 1 + ) + + for rank, row in enumerate(visual_rows): + key = (str(row.get("document_id") or ""), str(row.get("chunk_id") or "")) + rows_by_key.setdefault(key, dict(row)) + rrf_scores[key] = rrf_scores.get(key, 0.0) + CHANNEL_WEIGHT_VISUAL / ( + RRF_K + rank + 1 + ) + + ranked_keys = sorted( + rrf_scores, + key=lambda key: rrf_scores[key], + reverse=True, + )[:top_k] + + fused_rows: list[dict[str, Any]] = [] + for key in ranked_keys: + merged = dict(rows_by_key[key]) + merged["score"] = rrf_scores[key] + fused_rows.append(merged) + + if fused_rows: + normalize_row_scores( + fused_rows, + source_field="score", + target_field="discovery_score", + default=0.5, + ) + + return DiscoveryResult( + status="discovery_done", + payload={"fused_rows": fused_rows}, + latency_ms=int(lexical_result.latency_ms or 0) + + int(visual_result.latency_ms or 0), + ) diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 4a57bfe7c..8386de449 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -18,6 +18,7 @@ _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} _CLIENT_ARTIFACT_DIRS = {"images", "tables", "page_pdfs", "page_citation_assets"} _INTERNAL_RAW_FILES = {"source.pdf"} +_INTERNAL_RAW_PREFIXES = {"phase_b"} _MAX_RESULT_UPLOAD_CONCURRENCY = 64 _upload_executor_lock = threading.Lock() _upload_executors: dict[int, ThreadPoolExecutor] = {} @@ -387,6 +388,8 @@ def _normalize_upload_artifact_ref(self, artifact_ref: str | None) -> str | None return None if normalized in _INTERNAL_RAW_FILES: return normalized + if normalized.split("/", 1)[0] in _INTERNAL_RAW_PREFIXES: + return normalized return self.normalize_artifact_ref(normalized) def _is_excluded_file(self, file_name: str) -> bool: diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py index fca0df1b5..42161a4dc 100644 --- a/packages/shared-python/shared/testing/contract_runtime.py +++ b/packages/shared-python/shared/testing/contract_runtime.py @@ -55,6 +55,7 @@ "main", "custom_openapi", "shared.core.config", + "shared.core.config.ai", "shared.core.config.app", "shared.core.config.base", "shared.core.config.billing", diff --git a/packages/shared-python/shared/tests/test_agent_explore_prompt.py b/packages/shared-python/shared/tests/test_agent_explore_prompt.py index 256413552..0a19b4e2d 100644 --- a/packages/shared-python/shared/tests/test_agent_explore_prompt.py +++ b/packages/shared-python/shared/tests/test_agent_explore_prompt.py @@ -38,7 +38,7 @@ def test_loop_rules_state_round_read_and_pick_contract() -> None: assert "On your first\n turn you have no results yet, so do not call corpus.read." in LOOP_RULES assert "Calls in the same turn cannot use each other's results." in LOOP_RULES assert "the next turn is\n a pick phase: you can only call corpus.pick." in LOOP_RULES - assert "Body text you read but did not pick cannot be read again." in LOOP_RULES + assert "Body text you read (picked or not) cannot be read again." in LOOP_RULES assert "the steps used so far" in LOOP_RULES for removed in ("expire", "to pick later", '"pick" field', "remaining budget"): assert removed not in LOOP_RULES, removed diff --git a/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py b/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py index 6f153d9df..a8c09557a 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py +++ b/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py @@ -608,7 +608,7 @@ async def _discovery(*_args, **_kwargs): # noqa: ANN001 ) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + "shared.services.retrieval.agent_tools.tools.recall.hybrid_discovery", _discovery, ) result = await recall(explore_ctx, {"query": "intro"}) @@ -647,7 +647,7 @@ async def _discovery(*_args, **_kwargs): # noqa: ANN001 ) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + "shared.services.retrieval.agent_tools.tools.recall.hybrid_discovery", _discovery, ) result = await recall(explore_ctx, {"query": "dose"}) @@ -688,7 +688,7 @@ async def _discovery(*_args, **_kwargs): # noqa: ANN001 ) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + "shared.services.retrieval.agent_tools.tools.recall.hybrid_discovery", _discovery, ) monkeypatch.setattr( @@ -714,7 +714,7 @@ async def _discovery(*_args, **_kwargs): # noqa: ANN001 ) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + "shared.services.retrieval.agent_tools.tools.recall.hybrid_discovery", _discovery, ) monkeypatch.setattr( diff --git a/packages/shared-python/shared/tests/test_agent_tools_map_scope.py b/packages/shared-python/shared/tests/test_agent_tools_map_scope.py index 3a1b87199..efab78c34 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_map_scope.py +++ b/packages/shared-python/shared/tests/test_agent_tools_map_scope.py @@ -19,6 +19,7 @@ from sqlalchemy.orm import Session from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job import Job from shared.models.database.job_result import JobResult from shared.services.retrieval.agent_tools.registry import ( MAP_TOOL_CHAR_BUDGET, @@ -247,6 +248,7 @@ def map_ctx() -> ToolContext: def _disable_fk(dbapi_connection, _connection_record) -> None: # noqa: ANN001 dbapi_connection.execute("PRAGMA foreign_keys=OFF") + Job.__table__.create(engine) JobResult.__table__.create(engine) Document.__table__.create(engine) DocumentSection.__table__.create(engine) @@ -369,7 +371,7 @@ async def _discovery(*_args: object, **kwargs: object) -> DiscoveryResult: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + "shared.services.retrieval.agent_tools.tools.recall.hybrid_discovery", _discovery, ) result = await recall( diff --git a/packages/shared-python/shared/tests/test_embedding_gateway_client.py b/packages/shared-python/shared/tests/test_embedding_gateway_client.py new file mode 100644 index 000000000..699448094 --- /dev/null +++ b/packages/shared-python/shared/tests/test_embedding_gateway_client.py @@ -0,0 +1,208 @@ +"""I/O contract for Vercel Gemini Embedding 2 page-image and query vectors.""" + +from __future__ import annotations + +import os + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import httpx + +from shared.core.exceptions.domain_exceptions import LLMServiceException +from shared.services.ai.embedding import multimodal_client +from shared.services.ai.embedding.multimodal_client import ( + DEFAULT_EMBEDDING_DIM, + DEFAULT_EMBEDDING_MODEL, + _MAX_EMBEDDING_RETRIES, + _build_embedding_request, + _call_gateway_embedding, + _extract_embedding_vector, + embed_images, +) + + +def _ok_response() -> httpx.Response: + return httpx.Response( + 200, + json={ + "data": [{"embedding": [0.1] * DEFAULT_EMBEDDING_DIM}], + "usage": {"prompt_tokens": 1}, + }, + request=httpx.Request("POST", "https://ai-gateway.vercel.sh/v1/embeddings"), + ) + + +def _status_response(status_code: int, text: str = "no") -> httpx.Response: + return httpx.Response( + status_code, + text=text, + request=httpx.Request("POST", "https://ai-gateway.vercel.sh/v1/embeddings"), + ) + + +def _ssl_error() -> httpx.ReadError: + return httpx.ReadError( + "[SSL] record layer failure", + request=httpx.Request("POST", "https://ai-gateway.vercel.sh/v1/embeddings"), + ) + + +def _patch_gateway(monkeypatch: pytest.MonkeyPatch, post_fn) -> None: + class _Client: + def post(self, *args, **kwargs): + return post_fn(*args, **kwargs) + + monkeypatch.setattr(multimodal_client, "_resolve_vercel_api_key", lambda: "test-key") + monkeypatch.setattr(multimodal_client, "get_sync_client", lambda: _Client()) + monkeypatch.setattr(multimodal_client.time, "sleep", lambda _seconds: None) + + +def test_query_request_is_openai_text_input_without_dimensions() -> None: + body = _build_embedding_request( + model=DEFAULT_EMBEDDING_MODEL, + text="table of contents", + ) + assert body == { + "model": DEFAULT_EMBEDDING_MODEL, + "input": "table of contents", + } + + +def test_image_request_uses_gemini_inline_image_not_data_uri_text() -> None: + body = _build_embedding_request( + model=DEFAULT_EMBEDDING_MODEL, + image_mime="image/png", + image_bytes=b"\x89PNG", + ) + assert body["model"] == DEFAULT_EMBEDDING_MODEL + assert body["input"] == "" + assert "dimensions" not in body + content = body["providerOptions"]["google"]["content"] + part = content[0][0]["inlineData"] + assert part["mimeType"] == "image/png" + assert part["data"] + assert not str(body["input"]).startswith("data:") + + +def test_extract_vector_requires_openai_shape_and_3072_dims() -> None: + vector = _extract_embedding_vector( + {"data": [{"embedding": [0.1] * DEFAULT_EMBEDDING_DIM}]}, + expected_dim=DEFAULT_EMBEDDING_DIM, + model=DEFAULT_EMBEDDING_MODEL, + ) + assert len(vector) == DEFAULT_EMBEDDING_DIM + assert vector[0] == pytest.approx(0.1) + + with pytest.raises(LLMServiceException, match="dimension mismatch"): + _extract_embedding_vector( + {"data": [{"embedding": [0.1] * 1024}]}, + expected_dim=DEFAULT_EMBEDDING_DIM, + model=DEFAULT_EMBEDDING_MODEL, + ) + + with pytest.raises(LLMServiceException, match="no embeddings"): + _extract_embedding_vector( + {"output": {"embeddings": [{"embedding": [0.1] * DEFAULT_EMBEDDING_DIM}]}}, + expected_dim=DEFAULT_EMBEDDING_DIM, + model=DEFAULT_EMBEDDING_MODEL, + ) + + +def test_ssl_blip_retries_then_succeeds(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"n": 0} + + def post_fn(*_args, **_kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise _ssl_error() + return _ok_response() + + _patch_gateway(monkeypatch, post_fn) + vector = _call_gateway_embedding(text="query", seed_bytes=b"query") + assert len(vector) == DEFAULT_EMBEDDING_DIM + assert calls["n"] == 2 + + +def test_http_500_retries_then_succeeds(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"n": 0} + + def post_fn(*_args, **_kwargs): + calls["n"] += 1 + if calls["n"] == 1: + return _status_response(500, "upstream") + return _ok_response() + + _patch_gateway(monkeypatch, post_fn) + vector = _call_gateway_embedding(text="query", seed_bytes=b"query") + assert len(vector) == DEFAULT_EMBEDDING_DIM + assert calls["n"] == 2 + + +def test_http_400_does_not_retry(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"n": 0} + + def post_fn(*_args, **_kwargs): + calls["n"] += 1 + return _status_response(400, "bad request") + + _patch_gateway(monkeypatch, post_fn) + with pytest.raises(LLMServiceException, match="HTTP 400"): + _call_gateway_embedding(text="query", seed_bytes=b"query") + assert calls["n"] == 1 + + +def test_ssl_exhausted_retries_then_fails(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"n": 0} + + def post_fn(*_args, **_kwargs): + calls["n"] += 1 + raise _ssl_error() + + _patch_gateway(monkeypatch, post_fn) + with pytest.raises(LLMServiceException, match="after 3 attempts"): + _call_gateway_embedding(text="query", seed_bytes=b"query") + assert calls["n"] == _MAX_EMBEDDING_RETRIES + + +def test_embed_images_keeps_other_pages_after_one_ssl_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"n": 0} + + def post_fn(*_args, **_kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise _ssl_error() + return _ok_response() + + _patch_gateway(monkeypatch, post_fn) + vectors = embed_images([b"page-a", b"page-b"], concurrency=2) + assert len(vectors) == 2 + assert all(item is not None and len(item) == DEFAULT_EMBEDDING_DIM for item in vectors) + assert calls["n"] == 3 + + +def test_embed_images_skips_page_after_retries_and_keeps_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"n": 0} + + def post_fn(*_args, **_kwargs): + calls["n"] += 1 + if calls["n"] <= _MAX_EMBEDDING_RETRIES: + raise _ssl_error() + return _ok_response() + + _patch_gateway(monkeypatch, post_fn) + vectors = embed_images([b"page-a", b"page-b"], concurrency=1) + assert vectors[0] is None + assert vectors[1] is not None + assert len(vectors[1]) == DEFAULT_EMBEDDING_DIM + assert calls["n"] == _MAX_EMBEDDING_RETRIES + 1 diff --git a/pyproject.toml b/pyproject.toml index 165a07411..a875913ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,6 @@ exclude = [ "packages/shared-python/build", "packages/shared-python/dist", "packages/shared-python/shared/tests", - "deprecated", ] executionEnvironments = [ { root = "apps/api", extraPaths = ["packages/shared-python", "."] }, @@ -65,7 +64,6 @@ src = [ "apps/worker", "packages/shared-python", ] -exclude = ["deprecated"] [tool.ruff.lint] select = ["E", "F"] diff --git a/pytest.ini b/pytest.ini index cd7e47820..e5cdcc172 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,6 +1,5 @@ [pytest] -norecursedirs = deprecated .* -addopts = --ignore=deprecated +norecursedirs = .* asyncio_mode = auto asyncio_default_fixture_loop_scope = function filterwarnings = diff --git a/uv.lock b/uv.lock index 873fcd412..80697a626 100644 --- a/uv.lock +++ b/uv.lock @@ -1551,7 +1551,7 @@ requires-dist = [ { name = "kombu", specifier = "==5.5.4" }, { name = "loguru", specifier = "==0.7.3" }, { name = "openai", specifier = ">=1.0.0" }, - { name = "pgvector", specifier = "==0.2.4" }, + { name = "pgvector", specifier = ">=0.4.1" }, { name = "pillow", specifier = "==12.2.0" }, { name = "posthog", specifier = "==7.18.1" }, { name = "psutil", specifier = "==5.9.8" }, @@ -2599,13 +2599,11 @@ wheels = [ [[package]] name = "pgvector" -version = "0.2.4" +version = "0.5.1" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] +sdist = { url = "https://files.pythonhosted.org/packages/f8/23/96aa38899fbf8e103766db608d6e42acac269a96e08f3003fe9da3396fed/pgvector-0.5.1.tar.gz", hash = "sha256:94998a54b801b1075d623b8fa677fcb8210a7977b88f8e2203ab115c155af2e4", size = 35714, upload-time = "2026-10-09T01:50:22.779Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/70/4121568743eff331240def4d0b0e949f3cd36f440435a69f967ebd1f0bc6/pgvector-0.2.4-py2.py3-none-any.whl", hash = "sha256:548e1f88d3c7433020c1c177feddad2f36915c262852d621f9018fcafff6870b", size = 9564, upload-time = "2023-11-24T17:26:57.442Z" }, + { url = "https://files.pythonhosted.org/packages/a2/8d/a9c2a531da0ebb54b4a7174450e8534a39db112a141ae3a437de28420111/pgvector-0.5.1-py3-none-any.whl", hash = "sha256:ec5bcd5ffaefe6ecb2dcc9564ca921d284564b969183bc837a144604773af8ea", size = 31056, upload-time = "2026-10-09T01:50:21.614Z" }, ] [[package]]