diff --git a/AGENTS.md b/AGENTS.md index 9e68dc180..c9afaa80c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -102,13 +102,14 @@ flowchart TB subgraph RETRIEVE["⑤ Retrieval (shared)"] Query["GET /v1/retrieval/query"] --> Pipeline["run_retrieval_query"] - Pipeline --> Channels["3-Channel BM25 (path/content/term)"] - Pipeline --> Agentic["WorkflowOrchestrator (Planner + DAG)"] - Channels --> RRF["RRF Fusion"] - Agentic --> Hydrate["hydrate_paths_to_rows"] - RRF --> Rank["_rank_candidates_by_path"] - Hydrate --> Rank + Pipeline --> Classic["classic_topk / small_corpus (use_agentic=False)"] + Pipeline --> MapNav["mapnav checklist (default / use_agentic≠False)"] + Classic --> Channels["3-Channel BM25 (path/content/term)"] + Channels --> Rank["rank_retrieval_candidates"] + MapNav --> NavSnap["nav_snapshot + run_nav_episode"] + NavSnap --> Bridge["nav_bridge referenced_chunks"] Rank --> Assemble["assemble_retrieval_results"] + Bridge --> Assemble Assemble --> Results["Cited Evidence Results"] end ``` @@ -544,22 +545,19 @@ debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting. Core retrieval internals are grouped by ownership: -- `execution/`: request shaping, route selection, legacy route execution, and public response projection. -- `search/`: lexical channels, scoring, section filters, and candidate ranking. +- `execution/`: request shaping, route selection (classic / mapnav / small_corpus), and public response projection. +- `search/`: lexical channels, scoring, section filters, candidate ranking, and classic `bottom_discovery`. - `hydration/`: row/path/reference hydration, inline assets, and result assembly. +- `nav/` + `nav_*.py`: map-nav checklist episode (PLANNER / HARVEST / CONTROL). +- `trace/`: `DecisionTraceStep` mapping and `TraceRecorder`. - `graph/`: document graph publication/query support. - `stats/`: retrieval hit recording. -- `workflow/`: query planning, retrieve step execution, and wallet state. -- `agentic/core/`: agentic run types, token budgets, runtime config, and traces. -- `agentic/discovery/`: bottom discovery and document selection. -- `agentic/navigation/`: section-tree navigation, selection hydration, and asset tools. -- `agentic/evidence/`: evidence tree rendering and budget trimming. ### Two Retrieval Modes -The system supports two modes, controlled globally by `RETRIEVAL_AGENTIC_ENABLED` and locally via the per-request `use_agentic` toggle. +Per-request `use_agentic`: `False` → classic 3-channel top-K; `None`/`True` → map-nav (default). -#### Legacy Mode (3-Channel RRF) +#### Classic Mode (3-Channel RRF) ```mermaid flowchart LR @@ -569,40 +567,16 @@ flowchart LR P --> RRF["RRF Fusion (k=60)"] C --> RRF T --> RRF - RRF --> Graph[Legacy Graph Routing] - Graph --> Rank[Dual-priority ranking] + RRF --> Rank[rank_retrieval_candidates] Rank --> Assemble[hydration.result_assembly] ``` **Channel weights** (default): path=1.0, content=2.0, term=1.5 **RRF formula**: `score = weight / (k + rank + 1)` per channel, summed across channels. -#### Agentic Mode (Workflow Orchestrator) +#### Map-nav Mode (default) -The agentic pipeline uses `WorkflowOrchestrator` to handle complex queries via a DAG-based planning and budget-constrained execution engine: - -1. **Planning (`PlannerAgent`)**: The query is analyzed and decomposed into a DAG of retrieval steps. - - Simple queries generate a single `retrieve` step. - - Complex queries are broken into multiple `retrieve` steps. KNOWHERE does not plan answer synthesis steps. -2. **Budget Ledger (`BudgetLedger`)**: A strict token budget mechanism is enforced across the entire DAG execution. If the budget is exhausted, the pipeline halts safely and returns the best-effort evidence collected so far. -3. **Execution (`RetrievalAgent`)**: For each `retrieve` step, a multi-phase navigation engine runs: - - **Phase 1 (Discovery)**: 3-channel RRF keyword search and KG document selection. - - **Phase 2 (Navigation)**: Constrained Breadth-First Search (BFS) over the document's section tree. Discovered orphan leaves are merged into the tree to prevent data loss. - - **Phase 3 (Evidence Rendering)**: The hydrated document tree is rendered as `evidence_text`. -4. **Evidence-Only Contract**: Retrieval responses always expose `evidence_text` as the primary output. `answer_text` is retained only as a deprecated empty string. Downstream agents decide whether the evidence is sufficient and synthesize answers outside KNOWHERE. - -### Tree Rendering & Hydration - -Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is now determined dynamically by the `DocTreeNode` structure: -- **Structural Context (Outlines)**: Sections not drilled into are simply rendered as structural outlines (`title` + `summary`) to guide the LLM. -- **Leaf Content (Hydration)**: Sections that the LLM explicitly selects for drill-down have their raw chunks (`text`, `image`, `table`) fully hydrated into the `leaf_content` of the tree. -- **Multi-Modal Inline Embedding**: During hydration, connected inline assets (images/tables) are natively resolved and embedded directly into the text chunk content, supporting multi-modal LLM processing without brittle string-replacement placeholders. - -**`_rank_candidates_by_path()`** — Dual-priority ranking: - -- When agent results exist: agent_score is primary, discovery_score is tiebreaker -- Rows with agent_score=0 are demoted to fallback pool -- Sort key: `(agent_score, discovery_score, dual_hit_flag, importance_norm_score)` +Default agentic path is checklist map-nav (`nav/`): PLANNER (`plan_query`) → HARVEST (`execute_plan` / `harvest`, recursive DISPATCH) → CONTROL (`plan_control`). Episode config lives in `nav_config.py`. Exit bridge expands kept chunks to `referenced_chunks`; `decision_trace` is mapped in `trace/mapnav.py`. Token hard-stop uses `NavConfig.token_limit`. ### Result Assembly diff --git a/apps/api/.env.example b/apps/api/.env.example index e33e8bc37..400d74747 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -99,10 +99,10 @@ ARK_API_KEY= # IMAGE_MODEL_MAX=qwen3.6-flash # Optional retrieval overrides have code defaults. Retrieval is evidence-only: -# evidence_text is the primary output and answer_text is always empty. Set -# RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy -# 3-channel RRF mode. -# RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS=10.0 +# evidence_text is the primary output and answer_text is always empty. +# Default path is map-nav (PLANNER+HARVEST+CONTROL); set use_agentic=false for +# classic 3-channel RRF. Classic BM25 may use Postgres FTS prefilter: +# RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT=2000 # File handling defaults SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md,.html,.htm diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index bdba6660e..758f111e0 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -72,7 +72,10 @@ class RetrievalQueryRequest(BaseModel): ) use_agentic: bool | None = Field( None, - description="Set to true to enable agentic retrieval (LLM doc-select + navigation). Default (None/false) uses classic 3-channel top-K.", + description=( + "Map-nav (PLANNER+HARVEST+CONTROL) is the default when unset/true. " + "Set false to force classic 3-channel top-K retrieval." + ), ) @field_validator("channels") diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index 4223ccca7..81210b866 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -157,7 +157,7 @@ async def query_documents( top_k=top_k, exclude_document_ids=exclude_document_ids, exclude_sections=[item for item in exclude_sections], - use_agentic=True, + use_agentic=None, ) return to_mcp_query_response(response) diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index 270563328..a0f049c29 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -48,7 +48,7 @@ _PUBLIC_MODE_SELECTOR_FIELDS = {"mode", "processing"} _PARSE_TRACK_FIELD = "parse_track" _PAGE_MEMORY_FIELD_PREFIX = "page_memory" -_PUBLIC_COMPATIBILITY_EXTRA_FIELDS = frozenset({_PARSE_TRACK_FIELD}) +_PUBLIC_COMPATIBILITY_EXTRA_FIELDS = frozenset({_PARSE_TRACK_FIELD, "result_mode"}) IngestionCommandFactory = Callable[[str], DocumentIngestionCommand] diff --git a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py deleted file mode 100644 index 09858a1df..000000000 --- a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py +++ /dev/null @@ -1,315 +0,0 @@ -from shared.services.retrieval.agentic.navigation.actions import build_legal_actions -from shared.services.retrieval.agentic.navigation.state import RejectionRecord - - -def _rejected(paths: dict[str, str]) -> dict[str, RejectionRecord]: - """Build a rejection ledger from {path: reason}.""" - return { - path: RejectionRecord(path=path, reason=reason, step=1, detail="") - for path, reason in paths.items() - } - - -def test_discovery_hint_is_projected_as_collect_action() -> None: - action_set = build_legal_actions( - items=[], - current_scope=None, - collected_paths=[], - expanded_scopes=set(), - discovery_hints=[ - { - "section_path": "2 阶段性调整还是牛熊切换？ / 2.1 牛熊切换缘何开启?", - "discovery_score": 0.82, - "chunk_type": "text", - } - ], - rejected={}, - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - - assert len(action_set.collect) == 1 - action = action_set.collect[0] - assert action.id == "D1" - assert action.action == "COLLECT" - assert action.source == "discovery" - assert action.path == "2 阶段性调整还是牛熊切换？ / 2.1 牛熊切换缘何开启?" - assert action.score == 0.82 - - -def test_discovery_hint_under_collected_path_is_not_repeated() -> None: - action_set = build_legal_actions( - items=[], - current_scope=None, - collected_paths=[ - { - "path": "2 阶段性调整还是牛熊切换？", - "hydrate_mode": "chunks", - } - ], - expanded_scopes=set(), - discovery_hints=[ - { - "section_path": "2 阶段性调整还是牛熊切换？ / 2.1 牛熊切换缘何开启?", - "discovery_score": 0.82, - } - ], - rejected={}, - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - - assert action_set.collect == [] - - -def test_discovery_hint_under_tool_adjudicated_path_is_not_repeated() -> None: - """tool_adjudicated rejections are not revived by discovery this round.""" - action_set = build_legal_actions( - items=[], - current_scope=None, - collected_paths=[], - expanded_scopes=set(), - discovery_hints=[ - { - "section_path": "1、2016：机构行为助推行情演绎 / 二是英国“脱欧”影响下", - "discovery_score": 0.7, - } - ], - rejected=_rejected({ - "1、2016：机构行为助推行情演绎": "tool_adjudicated", - }), - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - - assert action_set.collect == [] - - -# ─── NavigationState ledger (Phase 0) ─────────────────────────────────────── - - -def test_mark_rejected_collect_records_tool_adjudicated_reason() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.mark_rejected_collect("Chapter 1", step=3, detail="no matching asset") - - assert "Chapter 1" in state.rejected - record = state.rejected["Chapter 1"] - assert record.reason == "tool_adjudicated" - assert record.step == 3 - assert record.detail == "no matching asset" - - -def test_mark_rejected_if_unproductive_records_navigational_abandon() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.mark_rejected_if_unproductive("Chapter 2", step=5, detail="back_from_unproductive") - - assert state.rejected["Chapter 2"].reason == "navigational_abandon" - - -def test_tool_adjudicated_overrides_weak_abandon_record() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.mark_rejected_if_unproductive("Chapter 3", step=2) - state.mark_rejected_collect("Chapter 3", step=4, detail="asset mismatch") - - # Stronger reason wins. - assert state.rejected["Chapter 3"].reason == "tool_adjudicated" - assert state.rejected["Chapter 3"].step == 4 - - -def test_weak_abandon_does_not_overwrite_strong_record() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.mark_rejected_collect("Chapter 4", step=1) - state.mark_rejected_if_unproductive("Chapter 4", step=5) - - assert state.rejected["Chapter 4"].reason == "tool_adjudicated" - assert state.rejected["Chapter 4"].step == 1 - - -def test_coverage_helpers_derive_from_collected_paths() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.add_collected( - {"path": "A", "hydrate_mode": "chunks", "confidence": 0.9}, - step=1, - scope_context=None, - ) - state.add_collected( - {"path": "B", "hydrate_mode": "outline", "confidence": 0.6}, - step=2, - scope_context=None, - ) - # A path upgraded from outline to full counts as covered, not outline. - state.add_collected( - {"path": "C", "hydrate_mode": "outline", "confidence": 0.5}, - step=3, - scope_context=None, - ) - state.add_collected( - {"path": "C", "hydrate_mode": "chunks", "confidence": 0.8}, - step=4, - scope_context=None, - ) - - assert state.covered_paths() == {"A", "C"} - assert state.outline_paths() == {"B"} - - -def test_snapshot_delta_records_rejection_reasons() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.step_count = 2 - state.mark_rejected_if_unproductive("X", step=2) - - state.step_count = 3 - state.mark_rejected_collect("Y", step=3, detail="asset mismatch") - - delta = state.snapshot_delta( - before_scope=None, - expanded_before=set(), - rejected_before={}, - collected_before_count=0, - ) - rejected_added = {item["path"]: item["reason"] for item in delta["rejected_added"]} - assert rejected_added == {"X": "navigational_abandon", "Y": "tool_adjudicated"} - - -def test_rejected_paths_with_reason_partitions_by_label() -> None: - from shared.services.retrieval.agentic.navigation.state import NavigationState - - state = NavigationState( - document_id="d1", - document_name="doc.pdf", - job_result_id="j1", - ) - state.mark_rejected_collect("A", step=1) - state.mark_rejected_if_unproductive("B", step=1) - state.mark_rejected_collect("C", step=1) - - assert state.rejected_paths_with_reason("tool_adjudicated") == {"A", "C"} - assert state.rejected_paths_with_reason("navigational_abandon") == {"B"} - - -# ─── Reason-aware action filtering (T7-style regression) ──────────────────── - - -def test_tool_adjudicated_rejection_blocks_collect_even_with_discovery() -> None: - """A tool-adjudicated path stays out of COLLECT even when discovery hints it.""" - action_set = build_legal_actions( - items=[], - current_scope=None, - collected_paths=[], - expanded_scopes=set(), - discovery_hints=[{"section_path": "X", "discovery_score": 0.95}], - rejected=_rejected({"X": "tool_adjudicated"}), - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - assert action_set.collect == [] - - -def test_navigational_abandon_is_revived_by_discovery_for_collect() -> None: - """A soft-abandoned path CAN still be COLLECTed when discovery signals it.""" - action_set = build_legal_actions( - items=[], - current_scope=None, - collected_paths=[], - expanded_scopes=set(), - discovery_hints=[{"section_path": "Y", "discovery_score": 0.9}], - rejected=_rejected({"Y": "navigational_abandon"}), - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - assert any(action.path == "Y" for action in action_set.collect) - - -def test_navigational_abandon_suppresses_expand_without_discovery_signal() -> None: - """EXPAND is suppressed for soft-abandoned scopes lacking any discovery signal.""" - items = [{"path": "Z", "level": 1, "is_leaf": False, "chunk_count": 5}] - action_set = build_legal_actions( - items=items, - current_scope=None, - collected_paths=[], - expanded_scopes=set(), - discovery_hints=[], - rejected=_rejected({"Z": "navigational_abandon"}), - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - assert any(action.path == "Z" for action in action_set.collect) - assert not any(action.path == "Z" for action in action_set.expand) - - -def test_navigational_abandon_revives_expand_with_discovery_signal() -> None: - """EXPAND is offered for soft-abandoned scopes when a discovery signal exists.""" - items = [{"path": "Z", "level": 1, "is_leaf": False, "chunk_count": 5}] - action_set = build_legal_actions( - items=items, - current_scope=None, - collected_paths=[], - expanded_scopes=set(), - discovery_hints=[{"section_path": "Z / child", "discovery_score": 0.7}], - rejected=_rejected({"Z": "navigational_abandon"}), - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - assert any(action.path == "Z" for action in action_set.expand) - - -def test_covered_path_excluded_from_actions() -> None: - """Regression: a path already collected as full evidence is not re-offered.""" - items = [{"path": "A", "level": 1, "is_leaf": False, "chunk_count": 3}] - action_set = build_legal_actions( - items=items, - current_scope=None, - collected_paths=[{"path": "A", "hydrate_mode": "chunks"}], - expanded_scopes=set(), - discovery_hints=[{"section_path": "A", "discovery_score": 0.9}], - rejected={}, - total_images=0, - total_tables=0, - budget_snapshot=None, - ) - assert not any(action.path == "A" for action in action_set.collect) - assert not any(action.path == "A" for action in action_set.expand) - diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py new file mode 100644 index 000000000..6b9583298 --- /dev/null +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -0,0 +1,199 @@ +"""Contract tests for the BM25 channel Postgres FTS prefilter. + +These run against a real Postgres so the prefilter is validated against the +same generated tsvector columns and GIN indexes production uses. A pure-Python +fake would not catch a mismatch between the query configuration and the one +the columns were generated with. +""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator + +import pytest +import pytest_asyncio +from shared.services.retrieval.search.channels import content_channel, path_channel +from shared.testing.contract_runtime import PostgreSQLProcess +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy import text + +_SCHEMA = """ +CREATE TABLE documents ( + document_id TEXT PRIMARY KEY, + user_id TEXT, + namespace TEXT, + status TEXT, + current_job_result_id INTEGER, + source_file_name TEXT +); +CREATE TABLE job_results (id INTEGER PRIMARY KEY, job_id TEXT); +CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); +CREATE TABLE document_chunks ( + id SERIAL PRIMARY KEY, + chunk_id TEXT, + document_id TEXT, + section_id TEXT, + chunk_type TEXT, + content TEXT, + source_chunk_path TEXT, + file_path TEXT, + chunk_metadata JSONB, + job_result_id INTEGER, + sort_order INTEGER, + content_search_text TEXT, + content_search_tsv TSVECTOR GENERATED ALWAYS AS + (to_tsvector('simple', COALESCE(content_search_text, ''))) STORED, + path_search_text TEXT, + path_search_tsv TSVECTOR GENERATED ALWAYS AS + (to_tsvector('simple', COALESCE(path_search_text, ''))) STORED, + term_search_text TEXT +); +CREATE INDEX idx_chunk_content_search_tsv ON document_chunks USING GIN (content_search_tsv); +CREATE INDEX idx_chunk_path_search_tsv ON document_chunks USING GIN (path_search_tsv); +""" + +_NOISE_ROWS = 300 + + +@pytest_asyncio.fixture +async def seeded_session( + postgresql_proc: PostgreSQLProcess, +) -> AsyncGenerator[AsyncSession, None]: + dsn = ( + f"postgresql+asyncpg://{postgresql_proc.user}@" + f"{postgresql_proc.host}:{postgresql_proc.port}/postgres" + ) + engine = create_async_engine(dsn, isolation_level="AUTOCOMMIT") + async with engine.begin() as conn: + await conn.execute(text("DROP SCHEMA IF EXISTS bm25_fts CASCADE")) + await conn.execute(text("CREATE SCHEMA bm25_fts")) + await conn.execute(text("SET search_path TO bm25_fts")) + for statement in filter(None, (s.strip() for s in _SCHEMA.split(";"))): + await conn.execute(text(statement)) + await conn.execute(text("INSERT INTO job_results VALUES (1, 'job1')")) + await conn.execute( + text( + "INSERT INTO documents VALUES " + "('d1', 'u1', 'ns1', 'active', 1, 'sample.pdf')" + ) + ) + await conn.execute(text("INSERT INTO document_sections VALUES ('s1', '/root')")) + await conn.execute( + text( + "INSERT INTO document_chunks " + "(chunk_id, document_id, section_id, chunk_type, content, " + " job_result_id, sort_order, content_search_text, path_search_text) " + "VALUES " + "('hit-en', 'd1', 's1', 'text', 'body', 1, 1, " + " 'alpha beta gamma', 'invoices alpha'), " + "('hit-cjk', 'd1', 's1', 'text', 'body', 1, 2, " + " '合同 条款 甲方', '合同 目录')" + ) + ) + await conn.execute( + text( + "INSERT INTO document_chunks " + "(chunk_id, document_id, section_id, chunk_type, content, " + " job_result_id, sort_order, content_search_text, path_search_text) " + "SELECT 'noise-' || i, 'd1', 's1', 'text', 'body', 1, i + 10, " + " 'filler unrelated wording ' || i, 'misc path ' || i " + "FROM generate_series(1, :noise) AS i" + ), + {"noise": _NOISE_ROWS}, + ) + + session_factory = async_sessionmaker(engine, expire_on_commit=False) + async with session_factory() as session: + await session.execute(text("SET search_path TO bm25_fts")) + yield session + await engine.dispose() + + +async def _content_hits(session: AsyncSession, query: str) -> list[str]: + rows = await content_channel( + session, + user_id="u1", + namespace="ns1", + query=query, + top_k=50, + exclude_document_ids=[], + exclude_sections=[], + ) + return [str(row["chunk_id"]) for row in rows] + + +@pytest.mark.asyncio +async def test_content_channel_returns_only_query_matching_chunks( + seeded_session: AsyncSession, +) -> None: + # The corpus holds hundreds of unrelated chunks. Before the prefilter every + # one of them was loaded into Python for BM25 scoring. + assert await _content_hits(seeded_session, "alpha") == ["hit-en"] + + +@pytest.mark.asyncio +async def test_content_channel_matches_cjk_tokens( + seeded_session: AsyncSession, +) -> None: + assert await _content_hits(seeded_session, "合同") == ["hit-cjk"] + + +@pytest.mark.asyncio +async def test_content_channel_uses_or_semantics_across_tokens( + seeded_session: AsyncSession, +) -> None: + # A row matching any single query token must survive, matching how the + # Python BM25 ranker admits rows. + hits = await _content_hits(seeded_session, "alpha 合同") + assert sorted(hits) == ["hit-cjk", "hit-en"] + + +@pytest.mark.asyncio +async def test_tsquery_operators_in_query_do_not_change_filter_shape( + seeded_session: AsyncSession, +) -> None: + # Tokens are lexed by Postgres as data. If operators leaked into tsquery + # syntax, "alpha & zzzz" would AND and drop the row. + assert await _content_hits(seeded_session, "alpha & zzzz") == ["hit-en"] + assert await _content_hits(seeded_session, "!alpha") == ["hit-en"] + + +@pytest.mark.asyncio +async def test_query_matching_nothing_returns_no_rows( + seeded_session: AsyncSession, +) -> None: + # The fallback re-runs the unfiltered scan, and BM25 then scores no row + # above zero, so the channel still yields nothing. + assert await _content_hits(seeded_session, "zzzznomatch") == [] + + +@pytest.mark.asyncio +async def test_path_channel_prefilters_on_path_search_tsv( + seeded_session: AsyncSession, +) -> None: + rows = await path_channel( + seeded_session, + user_id="u1", + namespace="ns1", + query="invoices", + top_k=50, + exclude_document_ids=[], + exclude_sections=[], + ) + assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] + + +@pytest.mark.asyncio +async def test_exclusions_still_apply_under_the_prefilter( + seeded_session: AsyncSession, +) -> None: + rows = await content_channel( + seeded_session, + user_id="u1", + namespace="ns1", + query="alpha", + top_k=50, + exclude_document_ids=["d1"], + exclude_sections=[], + ) + assert rows == [] diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index a1e704747..bb8efbbcb 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -294,6 +294,7 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( "namespace": "contract-demo", "query": "xAI investment", "top_k": 5, + "use_agentic": False, }, ) first_response = await api_client.post( @@ -316,6 +317,7 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( "namespace": "contract-demo", "query": "xAI investment", "top_k": 5, + "use_agentic": False, }, ) document_chunks_response = await api_client.get( diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 9b8eb1074..78dffed98 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -6,12 +6,8 @@ import pytest from httpx import AsyncClient from pytest import MonkeyPatch -from sqlalchemy.ext.asyncio import AsyncSession from tests.support.contract_database import ContractDatabase -from shared.services.retrieval.agentic.core.types import AgenticResult -from shared.services.retrieval.workflow.run_request import WorkflowRunRequest -from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult LLMFnInput = str | Sequence[dict[str, Any]] LLMFn = Callable[[LLMFnInput], Coroutine[Any, Any, str]] @@ -147,118 +143,7 @@ async def _seed_retrieval_chunk_for_existing_document( def _result_source(result: dict[str, object]) -> dict[str, object]: return cast(dict[str, object], result["source"]) -@pytest.mark.asyncio -async def test_agentic_workflow_should_pass_full_request_policy_to_step_adapter( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: MonkeyPatch, -) -> None: - captured_requests: list[dict[str, object]] = [] - - async def fake_plan( - self: object, - *, - query: str, - corpus_total_docs: int = 0, - corpus_total_chunks: int = 0, - ) -> QueryPlan: - return QueryPlan( - original_query=query, - steps=[PlannedStep(id="request-policy", sub_query="policy marker")], - final_strategy="concat_final_parts", - reasoning_summary=( - f"request policy contract for {corpus_total_docs} docs " - f"and {corpus_total_chunks} chunks" - ), - ) - - async def fake_retrieval_run( - self: object, - db: object, - **kwargs: object, - ) -> AgenticResult: - del self, db - captured_requests.append(kwargs) - return AgenticResult( - evidence_text="policy evidence", - answer_text="", - referenced_chunks=[ - { - "chunk_id": policy_document["chunk_id"], - "document_id": policy_document["document_id"], - "chunk_type": "page", - "section_path": policy_document["section_path"], - "file_path": "", - "job_id": policy_document["job_id"], - } - ], - router_used="contract_fake_agent", - ) - - async with developer_api_client_factory() as api_client: - policy_document = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-agentic-request-policy", - source_file_name="policy.pdf", - section_path="Root/Policy", - content="policy marker content", - chunk_type="page", - ) - await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-agentic-request-policy", - source_file_name="filler.pdf", - section_path="Root/Filler", - content="filler content", - chunk_type="page", - ) - monkeypatch.setattr( - "shared.services.retrieval.workflow.planner.QueryPlanner.plan", - fake_plan, - ) - monkeypatch.setattr( - "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", - fake_retrieval_run, - ) - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-agentic-request-policy", - "query": "policy marker", - "top_k": 1, - "chunk_types": ["page"], - "signal_paths": ["Root"], - "filter_mode": "keep", - "channels": ["content"], - "channel_weights": {"content": 2.0}, - "internal_recall_k": 23, - "threshold": 0.4, - "rerank": True, - "use_agentic": True, - }, - ) - - assert response.status_code == 200 - - assert len(captured_requests) == 1 - request = captured_requests[0] - assert request["user_id"] == "local-dev-user" - assert request["namespace"] == "contract-agentic-request-policy" - assert request["query"] == "policy marker" - assert request["top_k"] == 1 - assert request["exclude_document_ids"] == [] - assert request["exclude_sections"] == [] - assert request["chunk_types"] == {"page"} - assert request["signal_paths"] == ["Root"] - assert request["filter_mode"] == "keep" - assert request["channels"] == ["content"] - assert request["channel_weights"] == {"content": 2.0} - assert request["internal_recall_k"] == 23 - - -@pytest.mark.asyncio async def test_should_return_seeded_retrieval_results_for_the_authenticated_user( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] @@ -301,7 +186,6 @@ async def test_should_return_seeded_retrieval_results_for_the_authenticated_user } -@pytest.mark.asyncio async def test_should_default_the_namespace_to_default_when_it_is_omitted( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] @@ -334,7 +218,6 @@ async def test_should_default_the_namespace_to_default_when_it_is_omitted( assert _result_source(results[0])["document_id"] == seeded_document["document_id"] -@pytest.mark.asyncio async def test_should_return_empty_results_for_an_empty_query( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] @@ -357,7 +240,6 @@ async def test_should_return_empty_results_for_an_empty_query( assert response_json["referenced_chunks"] == [] -@pytest.mark.asyncio async def test_retrieval_should_use_classic_topk_when_agentic_is_false( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] @@ -394,305 +276,303 @@ async def test_retrieval_should_use_classic_topk_when_agentic_is_false( assert response_json["router_used"] == "classic_topk" -@pytest.mark.asyncio -async def test_agentic_retrieval_should_reference_root_only_document_content( +async def test_should_return_request_validation_failure_for_an_invalid_channel( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], - monkeypatch: MonkeyPatch, ) -> None: - monkeypatch.setenv("LLM_MOCK_ENABLED", "true") + async with developer_api_client_factory() as api_client: + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "default", + "query": "alpha", + "channels": ["invalid-channel"], + }, + ) + assert response.status_code == 400 + assert response.headers["x-request-id"] + + response_json = cast(dict[str, object], response.json()) + error = cast(dict[str, object], response_json["error"]) + details = cast(dict[str, object], error["details"]) + violations = cast(list[dict[str, object]], details["violations"]) + + assert response_json["success"] is False + assert error["code"] == "INVALID_ARGUMENT" + assert error["message"] == "Request validation failed" + assert violations[0]["field"] == "body.channels" + assert "Invalid channel" in cast(str, violations[0]["description"]) + + +async def test_should_exclude_matching_document_ids_from_the_response( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: async with developer_api_client_factory() as api_client: - rooted_document = await _seed_retrieval_document( + included_document = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-root-retrieval", - source_file_name="root-only.pdf", - section_path="Root", - content="root only diluted earnings marker content", + namespace="contract-retrieval", + source_file_name="included.pdf", + section_path="contract/included", + content="retrieval included content", ) - await _seed_retrieval_document( + excluded_document = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-root-retrieval", - source_file_name="filler.pdf", - section_path="filler/section", - content="unrelated filler content", + namespace="contract-retrieval", + source_file_name="excluded.pdf", + section_path="contract/excluded", + content="retrieval excluded content", ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-root-retrieval", - "query": "diluted earnings marker", - "top_k": 1, - "use_agentic": True, + "namespace": "contract-retrieval", + "query": "retrieval", + "exclude_document_ids": [excluded_document["document_id"]], }, ) assert response.status_code == 200 response_json = cast(dict[str, object], response.json()) - referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) results = cast(list[dict[str, object]], response_json["results"]) - assert response_json["router_used"] == "workflow_single_step" - assert { - "chunk_id": rooted_document["chunk_id"], - "document_id": rooted_document["document_id"], - "chunk_type": "text", - "section_path": "root-only.pdf", - "file_path": None, - "job_id": rooted_document["job_id"], - } in [ - {k: v for k, v in ref.items() if k != "score"} - for ref in referenced_chunks - ] - assert results[0]["content"] == "root only diluted earnings marker content" - assert results[0]["source"] == { - "document_id": rooted_document["document_id"], - "source_file_name": "root-only.pdf", - "section_path": "Root", - } + assert len(results) == 1 + assert _result_source(results[0])["document_id"] == included_document["document_id"] -@pytest.mark.asyncio -async def test_agentic_retrieval_should_reference_discovery_content_when_navigation_selects_nothing( +async def test_should_exclude_matching_sections_from_the_response( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], - monkeypatch: MonkeyPatch, ) -> None: - monkeypatch.setenv("LLM_MOCK_ENABLED", "true") - async with developer_api_client_factory() as api_client: - discovered_document = await _seed_retrieval_document( + included_document = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-discovery-fallback", - source_file_name="discovery.pdf", - section_path="Findings", - content="discovery fallback EBITDA margin marker content", + namespace="contract-retrieval", + source_file_name="included-section.pdf", + section_path="contract/keep", + content="section keep content", ) - await _seed_retrieval_document( + excluded_document = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-discovery-fallback", - source_file_name="filler.pdf", - section_path="filler/section", - content="unrelated filler content", + namespace="contract-retrieval", + source_file_name="excluded-section.pdf", + section_path="contract/exclude", + content="section exclude content", ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-discovery-fallback", - "query": "EBITDA margin marker", - "top_k": 1, - "use_agentic": True, + "namespace": "contract-retrieval", + "query": "section", + "exclude_sections": [ + { + "document_id": excluded_document["document_id"], + "section_path": excluded_document["section_path"], + } + ], }, ) assert response.status_code == 200 response_json = cast(dict[str, object], response.json()) - referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) results = cast(list[dict[str, object]], response_json["results"]) - assert response_json["router_used"] == "workflow_single_step" - assert { - "chunk_id": discovered_document["chunk_id"], - "document_id": discovered_document["document_id"], - "chunk_type": "text", - "section_path": discovered_document["section_path"], - "file_path": None, - "job_id": discovered_document["job_id"], - } in [ - {k: v for k, v in ref.items() if k != "score"} - for ref in referenced_chunks - ] - assert results[0]["content"] == "discovery fallback EBITDA margin marker content" - assert results[0]["source"] == { - "document_id": discovered_document["document_id"], - "source_file_name": "discovery.pdf", - "section_path": discovered_document["section_path"], - } + assert len(results) == 1 + assert _result_source(results[0])["document_id"] == included_document["document_id"] + assert _result_source(results[0])["section_path"] == included_document["section_path"] + + + +def _episode_keeping_chunks( + *, + documents: list[dict[str, str]], + evidence_text: str = "mapnav evidence", +) -> Any: + """Build a minimal EpisodeResult whose kept_chunks use real seeded chunk_ids.""" + from shared.services.retrieval.nav._compat import AgentStep, Chunk, EpisodeResult + + kept: list[Chunk] = [] + scored: list[tuple[Chunk, float]] = [] + for doc in documents: + chunk = Chunk( + node_id=doc["chunk_id"], + doc_id=doc["document_id"], + text=str(doc.get("content") or evidence_text), + line_ids=(0,), + section_id=doc.get("section_id"), + ) + kept.append(chunk) + scored.append((chunk, 1.0)) + return EpisodeResult( + representation="mapnav", + steps=[ + AgentStep( + step_idx=1, + action="query_plan", + detail={ + "plan": {"subgoals": [{"id": "s1"}], "coverage_checklist": []}, + "token_limit": 100000, + "tokens_used_total": 1, + "tokens_used_delta": 1, + "elapsed_ms": 1, + }, + ) + ], + scored_chunks=scored, + kept_chunks=kept, + evidence_text=evidence_text, + evidence_chars_actual=len(evidence_text), + retrieved_nodes=[d["chunk_id"] for d in documents], + stop_reason="completed", + ) + + +def _patch_run_nav_episode(monkeypatch: MonkeyPatch, episode: Any) -> None: + def _fake_run_nav_episode(*_args: Any, **_kwargs: Any) -> Any: + return episode + + monkeypatch.setattr( + "shared.services.retrieval.nav.run_nav_episode", + _fake_run_nav_episode, + ) @pytest.mark.asyncio -async def test_agentic_retrieval_should_not_send_table_artifacts_to_vlm( +async def test_mapnav_retrieval_should_return_seeded_chunk_via_fake_episode( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], monkeypatch: MonkeyPatch, ) -> None: - monkeypatch.setenv("LLM_MOCK_ENABLED", "true") - vlm_calls: list[LLMFnInput] = [] - - async def fake_vlm(prompt: LLMFnInput) -> str: - vlm_calls.append(prompt) - return '{"status":"DONE","answer":"unexpected table VLM answer"}' - - def fake_create_retrieval_vlm_fn(**_kwargs: object) -> LLMFn: - return fake_vlm - - class FakeResultStorage: - def generate_artifact_url( - self, - *, - job_id: str, - artifact_ref: str, - expires_in: int = 3600, - ) -> str | None: - del expires_in - return f"https://assets.example.com/{job_id}/{artifact_ref}" - - def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: - if not artifact_ref: - return None - normalized = artifact_ref.strip().replace("\\", "/").lstrip("/") - if not normalized: - return None - root_dir = normalized.split("/", 1)[0] - if root_dir not in {"images", "tables"}: - return None - return normalized - - def fake_get_result_storage() -> FakeResultStorage: - return FakeResultStorage() - async with developer_api_client_factory() as api_client: - monkeypatch.setattr( - "shared.services.retrieval.llm_adapter.create_retrieval_vlm_fn", - fake_create_retrieval_vlm_fn, - ) - monkeypatch.setattr( - "shared.services.retrieval.hydration.assets.get_result_storage", - fake_get_result_storage, - ) - table_document = await _seed_retrieval_document( + target = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-agentic-table-vlm-filter", - source_file_name="table-report.md", - section_path="Realdata Results Summary / Main Metrics", - content=( - "" - "" - "
budgetmetricvalue
1000Flat inspect_evidence_score_mean0.5674
" - ), - chunk_type="table", - file_path="tables/table-0-main-metrics.html", + namespace="contract-mapnav-seed", + source_file_name="target.pdf", + section_path="Findings", + content="mapnav seeded EBITDA marker content", ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-agentic-table-vlm-filter", - source_file_name="filler-table-report.md", - section_path="Appendix / Filler Metrics", - content=( - "" - "
metricvalue
unrelated filler metric999
" + namespace="contract-mapnav-seed", + source_file_name="filler.pdf", + section_path="filler/section", + content="unrelated filler content", + ) + target_with_content = {**target, "content": "mapnav seeded EBITDA marker content"} + _patch_run_nav_episode( + monkeypatch, + _episode_keeping_chunks( + documents=[target_with_content], + evidence_text="mapnav seeded EBITDA marker content", ), - chunk_type="table", - file_path="tables/table-1-filler-metrics.html", ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-agentic-table-vlm-filter", - "query": "budget 1000 Flat inspect_evidence_score_mean", + "namespace": "contract-mapnav-seed", + "query": "EBITDA marker", "top_k": 1, - "chunk_types": ["table"], "use_agentic": True, }, ) assert response.status_code == 200 - response_json = cast(dict[str, object], response.json()) referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) results = cast(list[dict[str, object]], response_json["results"]) - assert response_json["router_used"] == "workflow_single_step" - assert vlm_calls == [] - matching_references = [ - reference - for reference in referenced_chunks - if reference["chunk_id"] == table_document["chunk_id"] + assert response_json["router_used"] == "mapnav" + assert response_json["stop_reason"] == "completed" + assert isinstance(response_json.get("decision_trace"), list) + assert response_json["decision_trace"] + assert response_json["decision_trace"][-1]["phase"] == "terminal" + assert { + "chunk_id": target["chunk_id"], + "document_id": target["document_id"], + "chunk_type": "text", + "section_path": target["section_path"], + "file_path": "", + "job_id": target["job_id"], + } in [ + {k: v for k, v in ref.items() if k != "score"} + for ref in referenced_chunks ] - assert len(matching_references) == 1 - assert matching_references[0]["document_id"] == table_document["document_id"] - assert matching_references[0]["chunk_type"] == "table" - assert matching_references[0]["section_path"] == table_document["section_path"] - assert matching_references[0]["file_path"] == "tables/table-0-main-metrics.html" - assert matching_references[0]["job_id"] == table_document["job_id"] - assert str(matching_references[0]["asset_url"]).startswith( - "https://assets.example.com/" - ) - assert len(results) == 1 - assert results[0]["chunk_type"] == "table" - assert _result_source(results[0])["document_id"] == table_document["document_id"] + assert results[0]["content"] == "mapnav seeded EBITDA marker content" + assert results[0]["source"] == { + "document_id": target["document_id"], + "source_file_name": "target.pdf", + "section_path": target["section_path"], + } @pytest.mark.asyncio -async def test_agentic_retrieval_should_not_hydrate_references_outside_request_scope( +async def test_mapnav_retrieval_should_not_hydrate_references_outside_request_scope( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], monkeypatch: MonkeyPatch, ) -> None: - class FakeWorkflowOrchestrator: - async def run_request( - self, - _db: AsyncSession, - *, - request: WorkflowRunRequest, - ) -> WorkflowResult: - return WorkflowResult( - namespace=request.namespace, - query=request.query, - router_used="workflow_single_step", - answer_text="", - referenced_chunks=[ - { - "chunk_id": foreign_document["chunk_id"], - "document_id": foreign_document["document_id"], - "chunk_type": "text", - "section_path": foreign_document["section_path"], - "file_path": None, - "job_id": foreign_document["job_id"], - } - ], - ) - async with developer_api_client_factory() as api_client: - request_document = await _seed_retrieval_document( + await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-visible-scope", + namespace="contract-mapnav-visible", source_file_name="visible.pdf", section_path="visible/section", content="visible scoped content", ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-visible-scope", + namespace="contract-mapnav-visible", source_file_name="visible-filler.pdf", section_path="visible/filler", content="visible scoped filler content", ) - foreign_document = await _seed_retrieval_document( + foreign = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-foreign-scope", + namespace="contract-mapnav-foreign", source_file_name="foreign.pdf", section_path="foreign/section", content="foreign scoped content should not leak", ) + + def _fake_bridge(_episode: Any, _snapshot: Any) -> tuple[list[dict[str, Any]], dict[str, float]]: + return ( + [ + { + "chunk_id": foreign["chunk_id"], + "document_id": foreign["document_id"], + "chunk_type": "text", + "section_path": foreign["section_path"], + "file_path": None, + "job_id": foreign["job_id"], + } + ], + {foreign["chunk_id"]: 1.0}, + ) + + _patch_run_nav_episode( + monkeypatch, + _episode_keeping_chunks(documents=[{**foreign, "content": "x"}]), + ) monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", - FakeWorkflowOrchestrator, + "shared.services.retrieval.nav_bridge.build_referenced_chunks", + _fake_bridge, ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-visible-scope", + "namespace": "contract-mapnav-visible", "query": "visible", "top_k": 1, "use_agentic": True, @@ -700,71 +580,63 @@ async def run_request( ) assert response.status_code == 200 - response_json = cast(dict[str, object], response.json()) - referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) - results = cast(list[dict[str, object]], response_json["results"]) - - assert request_document["document_id"] != foreign_document["document_id"] - assert referenced_chunks == [] - assert results == [] + assert response_json["router_used"] == "mapnav" + assert response_json["referenced_chunks"] == [] + assert response_json["results"] == [] @pytest.mark.asyncio -async def test_agentic_retrieval_should_drop_references_that_do_not_match_the_hydrated_section( +async def test_mapnav_retrieval_should_drop_references_with_mismatched_section_path( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], monkeypatch: MonkeyPatch, ) -> None: - class FakeWorkflowOrchestrator: - async def run_request( - self, - _db: AsyncSession, - *, - request: WorkflowRunRequest, - ) -> WorkflowResult: - return WorkflowResult( - namespace=request.namespace, - query=request.query, - router_used="workflow_single_step", - answer_text="", - referenced_chunks=[ - { - "chunk_id": visible_chunk["chunk_id"], - "document_id": visible_chunk["document_id"], - "chunk_type": "text", - "section_path": "wrong/section", - "file_path": None, - "job_id": visible_chunk["job_id"], - } - ], - ) - async with developer_api_client_factory() as api_client: - visible_chunk = await _seed_retrieval_document( + visible = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-reference-section-match", + namespace="contract-mapnav-section-mismatch", source_file_name="visible.pdf", - section_path="right/section", + section_path="visible/section", content="visible scoped content", ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-reference-section-match", + namespace="contract-mapnav-section-mismatch", source_file_name="filler.pdf", section_path="filler/section", content="filler content", ) + + def _fake_bridge(_episode: Any, _snapshot: Any) -> tuple[list[dict[str, Any]], dict[str, float]]: + return ( + [ + { + "chunk_id": visible["chunk_id"], + "document_id": visible["document_id"], + "chunk_type": "text", + "section_path": "wrong/section/path", + "file_path": None, + "job_id": visible["job_id"], + } + ], + {visible["chunk_id"]: 1.0}, + ) + + _patch_run_nav_episode( + monkeypatch, + _episode_keeping_chunks(documents=[{**visible, "content": "x"}]), + ) monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", - FakeWorkflowOrchestrator, + "shared.services.retrieval.nav_bridge.build_referenced_chunks", + _fake_bridge, ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-reference-section-match", + "namespace": "contract-mapnav-section-mismatch", "query": "visible", "top_k": 1, "use_agentic": True, @@ -772,69 +644,44 @@ async def run_request( ) assert response.status_code == 200 - response_json = cast(dict[str, object], response.json()) - referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) - results = cast(list[dict[str, object]], response_json["results"]) - - assert referenced_chunks == [] - assert results == [] + assert response_json["router_used"] == "mapnav" + assert response_json["referenced_chunks"] == [] + assert response_json["results"] == [] @pytest.mark.asyncio -async def test_agentic_retrieval_should_fail_when_final_hydration_db_fails( +async def test_mapnav_retrieval_should_fail_when_final_hydration_db_fails( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], monkeypatch: MonkeyPatch, ) -> None: - class FakeWorkflowOrchestrator: - async def run_request( - self, - _db: AsyncSession, - *, - request: WorkflowRunRequest, - ) -> WorkflowResult: - return WorkflowResult( - namespace=request.namespace, - query=request.query, - router_used="workflow_single_step", - answer_text="", - referenced_chunks=[ - { - "chunk_id": visible_document["chunk_id"], - "document_id": visible_document["document_id"], - "chunk_type": "text", - "section_path": visible_document["section_path"], - "file_path": None, - "job_id": visible_document["job_id"], - } - ], - ) - async def fail_final_hydration(**_kwargs: object) -> object: raise RuntimeError("forced final hydration database failure") async with developer_api_client_factory() as api_client: - visible_document = await _seed_retrieval_document( + visible = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-final-hydration-failure", + namespace="contract-mapnav-hydration-failure", source_file_name="visible.pdf", section_path="visible/section", content="visible scoped content", ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-final-hydration-failure", + namespace="contract-mapnav-hydration-failure", source_file_name="filler.pdf", section_path="filler/section", content="filler content", ) from shared.services.retrieval.execution import routes as retrieval_routes - monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", - FakeWorkflowOrchestrator, + _patch_run_nav_episode( + monkeypatch, + _episode_keeping_chunks( + documents=[{**visible, "content": "visible scoped content"}] + ), ) monkeypatch.setattr( retrieval_routes, @@ -849,7 +696,7 @@ async def fail_final_hydration(**_kwargs: object) -> object: await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-final-hydration-failure", + "namespace": "contract-mapnav-hydration-failure", "query": "visible", "top_k": 1, "use_agentic": True, @@ -858,7 +705,7 @@ async def fail_final_hydration(**_kwargs: object) -> object: @pytest.mark.asyncio -async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_documents( +async def test_mapnav_should_preserve_same_chunk_id_across_documents( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -866,331 +713,61 @@ async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_i ) -> None: shared_chunk_id = f"chunk_{uuid4().hex[:12]}" - async def fake_plan( - self: object, - *, - query: str, - corpus_total_docs: int = 0, - corpus_total_chunks: int = 0, - ) -> QueryPlan: - return QueryPlan( - original_query=query, - steps=[ - PlannedStep(id="first", sub_query="first shared reference"), - PlannedStep(id="second", sub_query="second shared reference"), - ], - final_strategy="concat_final_parts", - reasoning_summary=( - f"forced two-step contract plan for {corpus_total_docs} docs " - f"and {corpus_total_chunks} chunks" - ), - ) - - async def fake_retrieval_run( - self: object, - db: object, - **kwargs: object, - ) -> AgenticResult: - query = str(kwargs["query"]) - document = first_document if query == "first shared reference" else second_document - return AgenticResult( - evidence_text=f"evidence for {document['document_id']}", - answer_text="", - referenced_chunks=[ - { - "chunk_id": shared_chunk_id, - "document_id": document["document_id"], - "chunk_type": "text", - "section_path": document["section_path"], - "file_path": "", - "job_id": document["job_id"], - } - ], - router_used="contract_fake_agent", - ) - async with developer_api_client_factory() as api_client: - first_document = await _seed_retrieval_document( + first = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-shared-chunk-id", + namespace="contract-mapnav-shared-chunk", source_file_name="first.pdf", - section_path="first/section", - content="shared deterministic content", + section_path="shared/first", + content="first shared reference content", chunk_id=shared_chunk_id, ) - second_document = await _seed_retrieval_document( + second_doc = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-shared-chunk-id", + namespace="contract-mapnav-shared-chunk", source_file_name="second.pdf", - section_path="second/section", - content="shared deterministic content", - chunk_id=shared_chunk_id, + section_path="shared/second-host", + content="host content for second document", ) - monkeypatch.setattr( - "shared.services.retrieval.workflow.planner.QueryPlanner.plan", - fake_plan, - ) - monkeypatch.setattr( - "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", - fake_retrieval_run, - ) - - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-shared-chunk-id", - "query": "show both shared references", - "top_k": 1, - "use_agentic": True, - }, - ) - - assert response.status_code == 200 - - response_json = cast(dict[str, object], response.json()) - referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) - results = cast(list[dict[str, object]], response_json["results"]) - - referenced_document_ids = { - cast(str, reference["document_id"]) for reference in referenced_chunks - } - result_document_ids = { - cast(str, _result_source(result)["document_id"]) for result in results - } - - assert referenced_document_ids == { - first_document["document_id"], - second_document["document_id"], - } - assert result_document_ids == { - first_document["document_id"], - second_document["document_id"], - } - - -@pytest.mark.asyncio -async def test_agentic_workflow_should_preserve_references_with_the_same_chunk_id_across_sections( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], - monkeypatch: MonkeyPatch, -) -> None: - shared_chunk_id = f"chunk_{uuid4().hex[:12]}" - - async def fake_plan( - self: object, - *, - query: str, - corpus_total_docs: int = 0, - corpus_total_chunks: int = 0, - ) -> QueryPlan: - return QueryPlan( - original_query=query, - steps=[ - PlannedStep(id="first", sub_query="first shared section"), - PlannedStep(id="second", sub_query="second shared section"), - ], - final_strategy="concat_final_parts", - reasoning_summary=( - f"forced section identity contract plan for {corpus_total_docs} docs " - f"and {corpus_total_chunks} chunks" - ), - ) - - async def fake_retrieval_run( - self: object, - db: object, - **kwargs: object, - ) -> AgenticResult: - query = str(kwargs["query"]) - chunk = first_chunk if query == "first shared section" else second_chunk - return AgenticResult( - evidence_text=f"evidence for {chunk['section_path']}", - answer_text="", - referenced_chunks=[ - { - "chunk_id": shared_chunk_id, - "document_id": chunk["document_id"], - "chunk_type": "text", - "section_path": chunk["section_path"], - "file_path": "", - "job_id": chunk["job_id"], - } - ], - router_used="contract_fake_agent", - ) - - async with developer_api_client_factory() as api_client: - first_chunk = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-shared-section-chunk-id", - source_file_name="same-document.pdf", - section_path="first/section", - content="repeated deterministic content", - chunk_id=shared_chunk_id, - ) - second_chunk = await _seed_retrieval_chunk_for_existing_document( + second = await _seed_retrieval_chunk_for_existing_document( user_id="local-dev-user", - namespace="contract-shared-section-chunk-id", - document=first_chunk, - section_path="second/section", - content="repeated deterministic content", + namespace="contract-mapnav-shared-chunk", + document=second_doc, + section_path="shared/second", + content="second shared reference content", chunk_id=shared_chunk_id, ) - monkeypatch.setattr( - "shared.services.retrieval.workflow.planner.QueryPlanner.plan", - fake_plan, - ) - monkeypatch.setattr( - "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", - fake_retrieval_run, + + _patch_run_nav_episode( + monkeypatch, + _episode_keeping_chunks( + documents=[ + {**first, "content": "first shared reference content"}, + {**second, "content": "second shared reference content"}, + ] + ), ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-shared-section-chunk-id", - "query": "show both shared section references", + "namespace": "contract-mapnav-shared-chunk", + "query": "shared reference", "top_k": 1, "use_agentic": True, }, ) assert response.status_code == 200 - response_json = cast(dict[str, object], response.json()) referenced_chunks = cast(list[dict[str, object]], response_json["referenced_chunks"]) results = cast(list[dict[str, object]], response_json["results"]) - referenced_section_paths = { - cast(str, reference["section_path"]) for reference in referenced_chunks - } - result_section_paths = { - cast(str, _result_source(result)["section_path"]) for result in results + assert response_json["router_used"] == "mapnav" + assert len(referenced_chunks) == 2 + assert {ref["document_id"] for ref in referenced_chunks} == { + first["document_id"], + second["document_id"], } - - assert referenced_section_paths == { - first_chunk["section_path"], - second_chunk["section_path"], - } - assert result_section_paths == { - first_chunk["section_path"], - second_chunk["section_path"], - } - - -@pytest.mark.asyncio -async def test_should_return_request_validation_failure_for_an_invalid_channel( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - async with developer_api_client_factory() as api_client: - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "default", - "query": "alpha", - "channels": ["invalid-channel"], - }, - ) - - assert response.status_code == 400 - assert response.headers["x-request-id"] - - response_json = cast(dict[str, object], response.json()) - error = cast(dict[str, object], response_json["error"]) - details = cast(dict[str, object], error["details"]) - violations = cast(list[dict[str, object]], details["violations"]) - - assert response_json["success"] is False - assert error["code"] == "INVALID_ARGUMENT" - assert error["message"] == "Request validation failed" - assert violations[0]["field"] == "body.channels" - assert "Invalid channel" in cast(str, violations[0]["description"]) - - -@pytest.mark.asyncio -async def test_should_exclude_matching_document_ids_from_the_response( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - async with developer_api_client_factory() as api_client: - included_document = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-retrieval", - source_file_name="included.pdf", - section_path="contract/included", - content="retrieval included content", - ) - excluded_document = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-retrieval", - source_file_name="excluded.pdf", - section_path="contract/excluded", - content="retrieval excluded content", - ) - - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-retrieval", - "query": "retrieval", - "exclude_document_ids": [excluded_document["document_id"]], - }, - ) - - assert response.status_code == 200 - - response_json = cast(dict[str, object], response.json()) - results = cast(list[dict[str, object]], response_json["results"]) - - assert len(results) == 1 - assert _result_source(results[0])["document_id"] == included_document["document_id"] - - -@pytest.mark.asyncio -async def test_should_exclude_matching_sections_from_the_response( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - async with developer_api_client_factory() as api_client: - included_document = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-retrieval", - source_file_name="included-section.pdf", - section_path="contract/keep", - content="section keep content", - ) - excluded_document = await _seed_retrieval_document( - user_id="local-dev-user", - namespace="contract-retrieval", - source_file_name="excluded-section.pdf", - section_path="contract/exclude", - content="section exclude content", - ) - - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "contract-retrieval", - "query": "section", - "exclude_sections": [ - { - "document_id": excluded_document["document_id"], - "section_path": excluded_document["section_path"], - } - ], - }, - ) - - assert response.status_code == 200 - - response_json = cast(dict[str, object], response.json()) - results = cast(list[dict[str, object]], response_json["results"]) - - assert len(results) == 1 - assert _result_source(results[0])["document_id"] == included_document["document_id"] - assert _result_source(results[0])["section_path"] == included_document["section_path"] + assert {ref["chunk_id"] for ref in referenced_chunks} == {shared_chunk_id} + assert len(results) == 2 diff --git a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py new file mode 100644 index 000000000..ca271f66f --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py @@ -0,0 +1,224 @@ +"""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.nav.nav_knowhere 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, + ) -> 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 + 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, + ) -> list[RouteRow]: + assert db is fresh_db + assert exclude_document_ids == [] + assert exclude_sections == [] + assert allowed_chunk_types 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", + channels=None, + channel_weights=None, + 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]["citation"]["document_id"] == "doc_contract" + assert outcome.completion_label == "MAPNAV RETRIEVAL" diff --git a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py b/apps/api/tests/contract/test_retrieval_workflow_session_contract.py deleted file mode 100644 index cbe7e472c..000000000 --- a/apps/api/tests/contract/test_retrieval_workflow_session_contract.py +++ /dev/null @@ -1,373 +0,0 @@ -from __future__ import annotations - -import asyncio -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import cast - -import pytest -from pytest import MonkeyPatch -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.exceptions.domain_exceptions import LLMServiceException -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.llm_adapter import LLMFn, create_retrieval_planner_fn -from shared.services.retrieval.workflow.orchestrator import DbSessionFactory, WorkflowOrchestrator -from shared.services.retrieval.workflow.plan_service import WorkflowPlanService -from shared.services.retrieval.workflow.planner import QueryPlanner -from shared.services.retrieval.workflow.run_request import WorkflowRunRequest -from shared.services.retrieval.workflow.step_runner import WorkflowStepRunner -from shared.services.retrieval.workflow.types import ( - PlannedStep, - QueryPlan, - StepResult, - WorkflowResult, -) - -RouteRow = dict[str, object] - - -def _build_planner( - llm_fn: LLMFn, - *, - timeout_seconds: float = 1.0, -) -> QueryPlanner: - return QueryPlanner( - llm_fn=llm_fn, - planner_ledger=None, - max_steps=3, - total_budget=100, - per_step_budget=10, - timeout_seconds=timeout_seconds, - ) - - -@pytest.mark.asyncio -async def test_workflow_planner_timeout_should_return_single_step_fallback() -> None: - async def slow_llm(_prompt: object) -> str: - await asyncio.sleep(0.05) - return "{}" - - planner = _build_planner(slow_llm, timeout_seconds=0.001) - - plan = await planner.plan(query="original query") - - assert plan.planner_status == "fallback" - assert plan.steps[0].sub_query == "original query" - assert plan.planner_error is not None - assert "timed out" in plan.planner_error - - -@pytest.mark.asyncio -async def test_workflow_planner_provider_error_should_return_single_step_fallback() -> None: - async def failing_llm(_prompt: object) -> str: - raise LLMServiceException(internal_message="provider unavailable") - - planner = _build_planner(failing_llm) - - plan = await planner.plan(query="provider failure query") - - assert plan.planner_status == "fallback" - assert plan.steps[0].sub_query == "provider failure query" - assert plan.planner_error is not None - - -@pytest.mark.asyncio -async def test_workflow_planner_invalid_json_should_return_single_step_fallback() -> None: - async def invalid_llm(_prompt: object) -> str: - return "not json" - - planner = _build_planner(invalid_llm) - - plan = await planner.plan(query="invalid planner output") - - assert plan.planner_status == "fallback" - assert plan.steps[0].sub_query == "invalid planner output" - assert plan.planner_error is not None - assert "JSON" in plan.planner_error - - -@pytest.mark.asyncio -async def test_workflow_planner_unexpected_code_error_should_propagate() -> None: - async def buggy_llm(_prompt: object) -> str: - raise RuntimeError("unexpected planner bug") - - planner = _build_planner(buggy_llm) - - with pytest.raises(RuntimeError, match="unexpected planner bug"): - await planner.plan(query="buggy planner query") - - -@pytest.mark.asyncio -async def test_workflow_planner_llm_should_pass_timeout_to_provider_client( - monkeypatch: MonkeyPatch, -) -> None: - observed_timeouts: list[object] = [] - - class FakeClient: - def chat_completion_with_usage( - self, - _prompt: object, - **kwargs: object, - ) -> tuple[str, dict[str, int]]: - observed_timeouts.append(kwargs.get("timeout")) - return "{}", {"total_tokens": 1} - - def fake_build_client_for_channel( - *, - channel: str, - model: str, - ) -> tuple[FakeClient, str]: - assert channel == "text" - return FakeClient(), model - - monkeypatch.setattr( - "shared.services.retrieval.llm_adapter._has_llm_credentials", - lambda: True, - ) - monkeypatch.setattr( - "shared.services.retrieval.llm_adapter._build_client_for_channel", - fake_build_client_for_channel, - ) - - planner_llm = create_retrieval_planner_fn(timeout_seconds=2.1) - - assert planner_llm is not None - await planner_llm("timeout contract") - assert observed_timeouts == [3] - - -@pytest.mark.asyncio -async def test_workflow_inventory_session_should_close_before_planner_starts( - monkeypatch: MonkeyPatch, -) -> None: - events: list[str] = [] - inventory_db = object() - - @asynccontextmanager - async def fake_db_factory() -> AsyncGenerator[AsyncSession, None]: - events.append("inventory_open") - try: - yield cast(AsyncSession, inventory_db) - finally: - events.append("inventory_close") - - async def fake_load_budget_inventory( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], - ) -> tuple[int, int, dict[str, int]]: - assert db is inventory_db - assert user_id == "contract-user" - assert namespace == "contract-namespace" - assert exclude_document_ids == [] - events.append("inventory_loaded") - return 3, 2, {} - - class RecordingPlanService: - async def load_or_create(self, **kwargs: object) -> QueryPlan: - events.append("planner_start") - assert events == [ - "inventory_open", - "inventory_loaded", - "inventory_close", - "planner_start", - ] - return QueryPlan.single_step(str(kwargs["query"])) - - class RecordingStepRunner: - async def run_step(self, **kwargs: object) -> None: - step = cast(PlannedStep, kwargs["step"]) - results_by_id = cast(dict[str, StepResult], kwargs["results_by_id"]) - results_by_id[step.id] = StepResult( - step_id=step.id, - sub_query=step.sub_query, - step_kind=step.step_kind, - depends_on=step.depends_on, - output_role=step.output_role, - status="done", - answer_text="", - ) - - def create_step_runner( - db_factory: DbSessionFactory, - parent_run_id: str, - ) -> WorkflowStepRunner: - assert db_factory is fake_db_factory - assert parent_run_id - return cast(WorkflowStepRunner, RecordingStepRunner()) - - monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator._load_budget_inventory", - fake_load_budget_inventory, - ) - orchestrator = WorkflowOrchestrator( - db_factory=fake_db_factory, - plan_service=cast(WorkflowPlanService, RecordingPlanService()), - step_runner_factory=create_step_runner, - ) - - await orchestrator.run_request( - cast(AsyncSession, object()), - request=WorkflowRunRequest( - user_id="contract-user", - namespace="contract-namespace", - query="session lifetime", - top_k=1, - exclude_document_ids=[], - exclude_sections=[], - ), - llm_fn=None, - ) - - assert events[:4] == [ - "inventory_open", - "inventory_loaded", - "inventory_close", - "planner_start", - ] - - -@pytest.mark.asyncio -async def test_agentic_route_should_release_route_session_before_fresh_final_hydration( - monkeypatch: MonkeyPatch, -) -> None: - events: list[str] = [] - route_db = _RecordingRouteSession(events) - fresh_db = object() - - @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") - - class FakeWorkflowOrchestrator: - async def run_request( - self, - db: AsyncSession, - *, - request: WorkflowRunRequest, - ) -> WorkflowResult: - events.append("workflow_start") - assert db is route_db - assert events[:2] == ["route_rollback", "workflow_start"] - return WorkflowResult( - namespace=request.namespace, - query=request.query, - router_used="workflow_single_step", - answer_text="", - referenced_chunks=[ - { - "chunk_id": "chunk_contract", - "document_id": "doc_contract", - "chunk_type": "text", - "section_path": "contract/section", - "file_path": None, - "job_id": "job_contract", - } - ], - ) - - 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, - ) -> ResolvedWorkflowReferences: - assert db is fresh_db - assert user_id == "contract-user" - assert namespace == "contract-namespace" - assert score_by_chunk_id 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, - ) -> list[RouteRow]: - assert db is fresh_db - assert exclude_document_ids == [] - assert exclude_sections == [] - assert allowed_chunk_types is None - events.append("assemble_results") - return rows - - monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", - FakeWorkflowOrchestrator, - ) - 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_agentic_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", - channels=None, - channel_weights=None, - rerank=False, - threshold=0.0, - internal_recall_k=None, - effective_recall_k=3, - use_agentic=True, - ) - ) - - assert events == [ - "route_rollback", - "workflow_start", - "fresh_db_open", - "resolve_references", - "assemble_results", - "fresh_db_close", - ] - assert outcome.response["results"][0]["citation"]["document_id"] == "doc_contract" - - -class _RecordingRouteSession: - def __init__(self, events: list[str]) -> None: - self._events = events - - async def rollback(self) -> None: - self._events.append("route_rollback") diff --git a/apps/api/tests/unit/test_bm25_channel_tsquery.py b/apps/api/tests/unit/test_bm25_channel_tsquery.py new file mode 100644 index 000000000..aea2d9dfa --- /dev/null +++ b/apps/api/tests/unit/test_bm25_channel_tsquery.py @@ -0,0 +1,45 @@ +"""Unit tests for BM25 channel Postgres FTS prefilter token preparation.""" + +from __future__ import annotations + +from shared.services.retrieval.search.channels import ( + _MAX_FTS_QUERY_TOKENS, + _prepare_fts_tokens, +) + + +def test_returns_empty_for_no_tokens() -> None: + assert _prepare_fts_tokens([]) == [] + + +def test_keeps_token_order() -> None: + assert _prepare_fts_tokens(["alpha", "beta"]) == ["alpha", "beta"] + + +def test_strips_surrounding_whitespace() -> None: + assert _prepare_fts_tokens([" alpha ", "beta"]) == ["alpha", "beta"] + + +def test_drops_blank_tokens() -> None: + assert _prepare_fts_tokens(["", " ", "alpha"]) == ["alpha"] + + +def test_returns_empty_when_every_token_is_blank() -> None: + assert _prepare_fts_tokens(["", " "]) == [] + + +def test_caps_token_count() -> None: + tokens = [f"tok{index}" for index in range(_MAX_FTS_QUERY_TOKENS + 25)] + assert len(_prepare_fts_tokens(tokens)) == _MAX_FTS_QUERY_TOKENS + + +def test_preserves_cjk_tokens() -> None: + assert _prepare_fts_tokens(["合同", "条款"]) == ["合同", "条款"] + + +def test_passes_tsquery_operators_through_untouched() -> None: + # Tokens travel to Postgres as a text[] parameter and are lexed there, so + # operator characters are data rather than syntax. Nothing is escaped or + # dropped here. + raw = ["alpha' & 'zzzz", "!beta", "a|b"] + assert _prepare_fts_tokens(raw) == raw diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 7a682e776..3742d1e30 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -94,9 +94,10 @@ ARK_API_KEY= # IMAGE_MODEL_MAX=qwen3.6-flash # Optional retrieval overrides have code defaults. Retrieval is evidence-only: -# evidence_text is the primary output and answer_text is always empty. Set -# RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy -# 3-channel RRF mode. +# evidence_text is the primary output and answer_text is always empty. +# Default path is map-nav (PLANNER+HARVEST+CONTROL); set use_agentic=false for +# classic 3-channel RRF. Classic BM25 may use Postgres FTS prefilter: +# RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT=2000 # Required for specific features: billing and analytics BILLING_ENABLED=false diff --git a/apps/worker/app/services/common/file_utils.py b/apps/worker/app/services/common/file_utils.py index c266afc97..90eb7a1eb 100644 --- a/apps/worker/app/services/common/file_utils.py +++ b/apps/worker/app/services/common/file_utils.py @@ -4,6 +4,10 @@ import pandas as pd +# Shared cap for cosmetic asset filename stems (images/tables). Keeps OS +# basename limits safe while preserving a short debug-friendly context. +MAX_ASSET_FILE_NAME_CHARS = 80 + def clean_file(path_, mode="remove", cols=None): """ diff --git a/apps/worker/app/services/document_agent/executor/prompts.py b/apps/worker/app/services/document_agent/executor/prompts.py index 8c225d2cc..06fe0458a 100644 --- a/apps/worker/app/services/document_agent/executor/prompts.py +++ b/apps/worker/app/services/document_agent/executor/prompts.py @@ -1,15 +1,15 @@ """Prompts for executor reflexion.""" REFLEXION_INSTRUCTIONS = ( - "You are the executor of a document profiling agent. Decide the next action " - "from the blackboard facts and available tools. Return strict JSON with keys: " - "action (tool_call or verdict_now), rationale, optional tool_name/tool_args, " - "optional verdict {status, rationale}. Use inspect.pages when more visual " - "evidence is needed, grep.text when native-PDF text evidence is needed, " - "propose.shard_plan when evidence is sufficient to shard, validate.anatomy_map " - "after a shard plan exists, and verdict only after validation succeeds. If a " - "tool failed or validation is invalid, either gather targeted evidence and " - "retry the relevant tool or abort with a clear rationale." + "You are the executor of a document profiling agent. Decide the next tool " + "call from the blackboard facts and available tools. Return strict JSON with " + "keys: action (must be tool_call), rationale, tool_name, tool_args. " + "Use inspect.pages when more visual evidence is needed, grep.text when " + "native-PDF text evidence is needed, propose.shard_plan when evidence is " + "sufficient to shard, validate.anatomy_map after a shard plan exists, and " + "the verdict tool to finish: verdict(status=success) only after validation " + "succeeds, or verdict(status=abort, rationale=...) only when the document " + "cannot be profiled. Do not invent other finish actions." ) __all__ = ["REFLEXION_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/executor/react_loop.py b/apps/worker/app/services/document_agent/executor/react_loop.py index a890f702f..9917ac4fa 100644 --- a/apps/worker/app/services/document_agent/executor/react_loop.py +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -56,27 +56,49 @@ def _compact_blackboard(ctx: ToolContext) -> dict[str, Any]: } +def _coerce_legacy_finish(data: dict[str, Any]) -> ReflexionDecision: + """Map obsolete ``action=verdict_now`` into a real tool call. + + Finish belongs exclusively to the ``verdict`` tool. A bare ``verdict_now`` + without an explicit status is treated as ready-to-shard, never as abort. + """ + rationale = str(data.get("rationale") or "") + raw_verdict = data.get("verdict") + if isinstance(raw_verdict, dict): + status = str(raw_verdict.get("status") or "").strip().lower() + if status in {"success", "abort"}: + return ReflexionDecision( + action="tool_call", + rationale=rationale, + tool_name="verdict", + tool_args={ + "status": status, + "rationale": str( + raw_verdict.get("rationale") or rationale or status + ), + }, + ) + return ReflexionDecision( + action="tool_call", + rationale=rationale or "Legacy verdict_now without status; propose shard plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) + + def _parse_decision(raw: str) -> ReflexionDecision: data = json.loads(raw) - action = str(data.get("action") or "tool_call") - if action not in {"tool_call", "verdict_now"}: + action = str(data.get("action") or "tool_call").strip().lower() + if action == "verdict_now": + return _coerce_legacy_finish(data if isinstance(data, dict) else {}) + if action != "tool_call": action = "tool_call" - verdict = None - if isinstance(data.get("verdict"), dict): - verdict_data = data["verdict"] - status = str(verdict_data.get("status") or "abort") - if status not in {"success", "abort"}: - status = "abort" - verdict = AgentVerdict( - status=status, # type: ignore[arg-type] - rationale=str(verdict_data.get("rationale") or data.get("rationale") or ""), - ) return ReflexionDecision( - action=action, # type: ignore[arg-type] + action="tool_call", rationale=str(data.get("rationale") or ""), tool_name=data.get("tool_name"), tool_args=dict(data.get("tool_args") or {}), - verdict=verdict, + verdict=None, ) @@ -98,6 +120,15 @@ def run(self) -> ExecutorResult: for round_index in range(self.max_rounds): pending_recovery_verdict: AgentVerdict | None = None decision, result = self._next_decision(round_index) + if decision.action != "tool_call": + # Defensive: only tool_call is a legal executor step. + decision = ReflexionDecision( + action="tool_call", + rationale=decision.rationale + or "Non-tool executor action coerced to propose.shard_plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) self.ctx.blackboard.global_signals.setdefault("reflexion_decisions", []).append( decision.to_dict() ) @@ -111,14 +142,9 @@ def run(self) -> ExecutorResult: tool_args=decision.tool_args, ) - tool_name: str | None = None - tool_args: dict[str, Any] = {} - if decision.action == "verdict_now": - verdict = decision.verdict or AgentVerdict( - status="abort", - rationale=decision.rationale or "Executor stopped without verdict.", - ) - if verdict.status == "success" and not ( + tool_name, tool_args = self._resolve_tool_call(decision) + if tool_name == "verdict" and str(tool_args.get("status") or "") == "success": + if not ( self.ctx.blackboard.validation_report and self.ctx.blackboard.validation_report.get("valid") is True ): @@ -131,11 +157,6 @@ def run(self) -> ExecutorResult: tool_args={}, ) tool_name, tool_args = self._resolve_tool_call(decision) - else: - self.ctx.blackboard.verdict = verdict - return ExecutorResult(verdict=verdict, rounds=round_index + 1) - else: - tool_name, tool_args = self._resolve_tool_call(decision) if not tool_name: verdict = AgentVerdict( @@ -205,6 +226,22 @@ def _is_deterministic_mode(self) -> bool: def _next_decision(self, round_index: int) -> tuple[ReflexionDecision, ToolResult]: if round_index == 0 and self._initial_decision is not None: decision = self._initial_decision + if decision.action != "tool_call": + decision = ReflexionDecision( + action="tool_call", + rationale=decision.rationale + or "Initial non-tool decision coerced to propose.shard_plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) + elif not decision.tool_name: + decision = ReflexionDecision( + action="tool_call", + rationale=decision.rationale + or "Initial decision missing tool; propose shard plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) return decision, ToolResult(status="ok", payload=decision.to_dict()) model = self.ctx.settings.get("executor_model") or self.ctx.settings.get("model") if not model: @@ -224,9 +261,13 @@ def _next_decision(self, round_index: int) -> tuple[ReflexionDecision, ToolResul est = estimate_tokens(prompt) if not self.ctx.budget.try_reserve("plan", est): decision = ReflexionDecision( - action="verdict_now", + action="tool_call", rationale="Planner budget exhausted.", - verdict=AgentVerdict(status="abort", rationale="Planner budget exhausted."), + tool_name="verdict", + tool_args={ + "status": "abort", + "rationale": "Planner budget exhausted.", + }, ) return decision, ToolResult( status="ok", @@ -303,4 +344,3 @@ def _deterministic_decision(self) -> ReflexionDecision: tool_name="validate.anatomy_map", tool_args={}, ) - diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index d078f9d58..d422a55b8 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -10,7 +10,10 @@ PageKind = Literal["normal", "landscape"] TocFailureKind = Literal["none", "confirm_failed", "rejected_all", "degraded"] -ReflexionAction = Literal["tool_call", "verdict_now"] +# Executor loop steps are always tool calls. Profile success/abort is owned +# exclusively by the ``verdict`` tool (AgentVerdict.status), not by a separate +# ReflexionAction shortcut. +ReflexionAction = Literal["tool_call"] VerdictStatus = Literal["success", "abort"] @@ -28,8 +31,6 @@ class PageFeature: height: float has_asset: bool is_blank_like: bool - # PDF-space boxes for detected assets; None when none were extracted. - asset_bboxes: list[dict[str, Any]] | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) diff --git a/apps/worker/app/services/document_agent/persist/__init__.py b/apps/worker/app/services/document_agent/persist/__init__.py index a34439602..26c0d806a 100644 --- a/apps/worker/app/services/document_agent/persist/__init__.py +++ b/apps/worker/app/services/document_agent/persist/__init__.py @@ -1,8 +1,13 @@ """Persist anatomy map artifacts.""" from app.services.document_agent.tools.persist_anatomy_map import ( + DOC_PROFILE_FILENAME, build_anatomy_map, persist_anatomy_map, ) -__all__ = ["build_anatomy_map", "persist_anatomy_map"] +__all__ = [ + "DOC_PROFILE_FILENAME", + "build_anatomy_map", + "persist_anatomy_map", +] diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index 7e5e1ce01..acf656003 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -156,7 +156,11 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec header_y=header_y, footer_y=footer_y, ) - next_action = str(data.get("next_action") or "ready_to_shard") + next_action = str(data.get("next_action") or "ready_to_shard").strip().lower() + # Legacy models may still emit verdict_now; that is not a planner finish + # signal — fall through to ready_to_shard so the executor owns success/abort. + if next_action == "verdict_now": + next_action = "ready_to_shard" tool_name: str | None = None tool_args: dict[str, Any] = {} if next_action == "inspect_more": @@ -171,12 +175,6 @@ def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDec if query: tool_name = "grep.text" tool_args = {"query": query, "max_results": 20} - elif next_action == "verdict_now": - return profile, ReflexionDecision( - action="verdict_now", - rationale=profile.rationale, - verdict=None, - ) if tool_name: return profile, ReflexionDecision( action="tool_call", diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py index 4ca93ef8e..2c6c48145 100644 --- a/apps/worker/app/services/document_agent/planner/prompts.py +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -17,10 +17,12 @@ "footer_y is the highest footer line you observe (smallest y) when any " "footer is present, otherwise null. When both are set, require " "header_y < footer_y. " - "next_action must be one of inspect_more, grep_text, ready_to_shard, " - "verdict_now. Use inspect_more only when extra page screenshots are needed. " + "next_action must be one of inspect_more, grep_text, ready_to_shard. " + "Use inspect_more only when extra page screenshots are needed. " "Use grep_text only for native PDFs when a global text search would clarify " - "structure. Do not output a fixed step plan." + "structure. Use ready_to_shard when evidence is sufficient to propose shards. " + "Do not finish or abort the profile run from next_action; the executor owns " + "success/abort via the verdict tool. Do not output a fixed step plan." ) __all__ = ["PLANNER_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index 5815a4120..f33717875 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -651,8 +651,9 @@ def _entries_to_tree(entries: list[dict[str, Any]]) -> list[TitleNode]: stack: list[tuple[int, TitleNode]] = [] for entry in entries: - raw_title = str(entry.get("heading") or "").strip() - title = clean_toc_title(raw_title) or normalize_heading_text(raw_title) + # Keep original TOC heading (incl. numbering). Prefix stripping belongs + # only in text/compact match helpers used for null-page parents. + title = normalize_heading_text(str(entry.get("heading") or "")) level = _safe_int(entry.get("level")) or 1 if not title or len(title) < 2: continue diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py index a2de3d295..1480eb4d8 100644 --- a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -6,6 +6,7 @@ import json import os import time +from dataclasses import dataclass, field from pathlib import Path from typing import Any, cast @@ -20,6 +21,7 @@ ) from app.services.document_agent.registry import register_tool from app.services.document_agent.tools.vlm_toc_extractor import ( + TOC_VLM_MAX_TOKENS, vlm_entries_to_toc_hierarchies, ) from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( @@ -34,81 +36,150 @@ MAX_BOUNDARY_ROUNDS = 6 MAX_TOC_PAGES = BOUNDARY_STEP_PAGES * MAX_BOUNDARY_ROUNDS # 30 +_CONFIRM_STAGE = "toc_confirm" +_CONFIRM_TOKENS_PER_PAGE = 800 + +_CONFIRM_PROMPT = ( + "You are a document structure analysis expert. " + "Below are screenshot(s) of candidate pages extracted from a PDF. " + "These pages contained keywords such as 'Table of Contents' / 'Contents' " + "during a text scan.\n\n" + "For each page, determine whether it is truly the **start page** of a " + "Table of Contents (TOC).\n\n" + "Criteria for a real TOC page:\n" + "- Contains a list of section titles paired with page numbers\n" + "- Titles are connected to page numbers via dots, ellipses, or spaces\n" + "- Titles have a systematic numbering scheme (e.g. 1. / 1.1 / Chapter 1)\n\n" + "NOT a TOC page:\n" + "- Body text that casually mentions 'contents'\n" + "- A page with only a 'Contents' heading but body text below\n\n" + "Return strict JSON (no markdown fences):\n" + '{"pages": [{"page": , "is_toc_start": true/false, ' + '"reason": "brief reason"}]}' +) + # -- PyMuPDF workers (must be top-level for multiprocessing pickle) ------------ @worker -def _render_single_page_worker( - queue, pdf_path: str, page_num: int, output_path: str, dpi: int +def _render_expand_window_worker( + queue, + pdf_path: str, + pages: list[int], + output_dir: str, + dpi: int, + anchor_page: int, ) -> None: + """Render one Phase2 expand window in a single child process. + + Opens the PDF once and writes ``toc_a{anchor}_p{page}.png`` for each page. + """ import pymupdf # type: ignore[import] + results: list[dict[str, Any]] = [] + doc = None try: doc = pymupdf.open(pdf_path) - idx = page_num - 1 - if 0 <= idx < doc.page_count: - page = doc[idx] - mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) - pix = page.get_pixmap(matrix=mat) - pix.save(output_path) + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + for page_num in pages: + idx = page_num - 1 + if not (0 <= idx < doc.page_count): + continue + pix = doc[idx].get_pixmap(matrix=mat) + png_path = os.path.join( + output_dir, + f"toc_a{anchor_page}_p{page_num}.png", + ) + pix.save(png_path) + results.append({"page": page_num, "png_path": png_path}) finally: try: - doc.close() + if doc is not None: + doc.close() except Exception: pass gc.collect() - queue.put({"ok": True, "png_path": output_path}) + queue.put({"ok": True, "results": results}) # -- VLM helpers --------------------------------------------------------------- -def _vlm_confirm_anchors( - anchor_pages: list[TocAnchorPage], - model: str, - budget: Any | None = None, -) -> tuple[list[TocAnchorPage], bool, list[TocEvidence]]: - """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" - from shared.services.ai.llm_overrides import get_vision_client +def _iter_chunks(items: list[Any], size: int) -> list[list[Any]]: + if size <= 0: + raise ValueError("chunk size must be positive") + return [items[i : i + size] for i in range(0, len(items), size)] + + +def _parse_confirm_items(raw: str) -> list[dict[str, Any]]: + data = json.loads(raw) + if isinstance(data, dict): + items = data.get("pages") or data.get("results") or data.get("data") or [] + if not items and len(data) == 1: + items = list(data.values())[0] + elif isinstance(data, list): + items = data + else: + items = [] + return [item for item in items if isinstance(item, dict)] + + +def _evidence_from_confirm_items( + items: list[dict[str, Any]], +) -> tuple[set[int], dict[int, TocEvidence]]: + confirmed_pages: set[int] = set() + evidence_by_page: dict[int, TocEvidence] = {} + for item in items: + if "page" not in item: + continue + page = int(item["page"]) + is_toc_start = bool(item.get("is_toc_start")) + if is_toc_start: + confirmed_pages.add(page) + raw_confidence = item.get("confidence") + try: + confidence = ( + float(raw_confidence) + if raw_confidence is not None + else (0.95 if is_toc_start else 0.05) + ) + except (TypeError, ValueError): + confidence = 0.95 if is_toc_start else 0.05 + evidence_by_page[page] = TocEvidence( + page_index=page, + source="vlm", + confidence=max(0.0, min(1.0, confidence)), + reason=str(item.get("reason") or ""), + ) + return confirmed_pages, evidence_by_page - if not anchor_pages: - return [], False, [] +def _confirm_anchor_chunk( + chunk: list[TocAnchorPage], + *, + model: str, + budget: Any | None, +) -> tuple[set[int], dict[int, TocEvidence], bool]: + """Confirm one BOUNDARY_STEP_PAGES-sized anchor chunk. + + Returns ``(confirmed_pages, evidence_by_page, failed)``. + """ import base64 - # Build multi-image message + from shared.services.ai.llm_overrides import get_vision_client + + if not chunk: + return set(), {}, False + content_parts: list[dict[str, Any]] = [ - { - "type": "text", - "text": ( - "You are a document structure analysis expert. " - "Below are screenshot(s) of candidate pages extracted from a PDF. " - "These pages contained keywords such as 'Table of Contents' / 'Contents' " - "during a text scan.\n\n" - "For each page, determine whether it is truly the **start page** of a " - "Table of Contents (TOC).\n\n" - "Criteria for a real TOC page:\n" - "- Contains a list of section titles paired with page numbers\n" - "- Titles are connected to page numbers via dots, ellipses, or spaces\n" - "- Titles have a systematic numbering scheme (e.g. 1. / 1.1 / Chapter 1)\n\n" - "NOT a TOC page:\n" - "- Body text that casually mentions 'contents'\n" - "- A page with only a 'Contents' heading but body text below\n\n" - "Return a strict JSON array (no markdown fences):\n" - '[{"page": , "is_toc_start": true/false, "reason": "brief reason"}]' - ), - } + {"type": "text", "text": _CONFIRM_PROMPT}, ] - - for anchor in anchor_pages: + for anchor in chunk: with open(anchor.png_path, "rb") as f: img_b64 = base64.b64encode(f.read()).decode() content_parts.append( - { - "type": "text", - "text": f"\n--- Page {anchor.page} ---", - } + {"type": "text", "text": f"\n--- Page {anchor.page} ---"} ) content_parts.append( { @@ -118,20 +189,22 @@ def _vlm_confirm_anchors( ) messages = cast(Any, [{"role": "user", "content": content_parts}]) - est = estimate_tokens(str(content_parts[0]["text"])) + len(anchor_pages) * 800 - stage = "toc_confirm" - if budget and not budget.try_reserve("visual", est, stage=stage): - logger.warning("[extract.toc] insufficient visual budget for anchor confirmation") - return [], True, [] + est = estimate_tokens(_CONFIRM_PROMPT) + len(chunk) * _CONFIRM_TOKENS_PER_PAGE + if budget and not budget.try_reserve("visual", est, stage=_CONFIRM_STAGE): + logger.warning( + "[extract.toc] insufficient visual budget for confirm chunk pages={}", + [a.page for a in chunk], + ) + return set(), {}, True try: client, resolved_model = get_vision_client(requested_model=model) - model = resolved_model or model + resolved = resolved_model or model raw, usage = client.chat_completion_with_usage( messages=messages, - model=model, + model=resolved, temperature=0.1, - max_tokens=500, + max_tokens=TOC_VLM_MAX_TOKENS, response_format={"type": "json_object"}, usage_task="document_agent.toc_anchor_confirm", ) @@ -140,72 +213,361 @@ def _vlm_confirm_anchors( "visual", actual=usage.get("total_tokens", est), est=est, - stage=stage, + stage=_CONFIRM_STAGE, ) - data = json.loads(raw) - if isinstance(data, dict): - items = data.get("pages") or data.get("results") or data.get("data") or [] - if not items and len(data) == 1: - items = list(data.values())[0] - elif isinstance(data, list): - items = data + confirmed_pages, evidence_by_page = _evidence_from_confirm_items( + _parse_confirm_items(raw) + ) + return confirmed_pages, evidence_by_page, False + except Exception as exc: + if budget: + budget.refund("visual", est=est, stage=_CONFIRM_STAGE) + logger.warning( + "[extract.toc] VLM confirm chunk failed pages={}: {}", + [a.page for a in chunk], + exc, + ) + return set(), {}, True + + +def _vlm_confirm_anchors( + anchor_pages: list[TocAnchorPage], + model: str, + budget: Any | None = None, +) -> tuple[list[TocAnchorPage], bool, list[TocEvidence]]: + """Phase 1: confirm TOC starts in BOUNDARY_STEP_PAGES batches (concurrent).""" + if not anchor_pages: + return [], False, [] + + from gevent.pool import Pool as GeventPool + + chunks = _iter_chunks(anchor_pages, BOUNDARY_STEP_PAGES) + logger.info( + "[extract.toc] Phase 1 confirm: {} anchors → {} chunks (size={}, concurrency={})", + len(anchor_pages), + len(chunks), + BOUNDARY_STEP_PAGES, + min(BOUNDARY_STEP_PAGES, len(chunks)), + ) + + pool = GeventPool(size=min(BOUNDARY_STEP_PAGES, len(chunks))) + jobs = [ + pool.spawn(_confirm_anchor_chunk, chunk, model=model, budget=budget) + for chunk in chunks + ] + pool.join() + + confirmed_pages: set[int] = set() + evidence_by_page: dict[int, TocEvidence] = {} + chunk_failures = 0 + chunk_successes = 0 + for job in jobs: + try: + chunk_confirmed, chunk_evidence, failed = job.get() + except Exception as exc: + chunk_failures += 1 + logger.warning("[extract.toc] confirm greenlet failed: {}", exc) + continue + if failed: + chunk_failures += 1 + continue + chunk_successes += 1 + confirmed_pages.update(chunk_confirmed) + evidence_by_page.update(chunk_evidence) + + confirm_failed = chunk_successes == 0 and chunk_failures > 0 + confirmed = [a for a in anchor_pages if a.page in confirmed_pages] + rejected = [a.page for a in anchor_pages if a.page not in confirmed_pages] + evidence = [ + evidence_by_page.get( + a.page, + TocEvidence( + page_index=a.page, + source="vlm", + confidence=0.05, + reason=( + "confirm batch failed for this candidate" + if confirm_failed + else "VLM response omitted this candidate page" + ), + ), + ) + for a in anchor_pages + ] + logger.info( + "[extract.toc] Phase 1 done: confirmed={} rejected={} " + "chunk_ok={} chunk_fail={} confirm_failed={}", + len(confirmed), + rejected, + chunk_successes, + chunk_failures, + confirm_failed, + ) + return confirmed, confirm_failed, evidence + + +# -- Phase 2: expand + extract per confirmed start ----------------------------- + + +@dataclass +class _TocRegionResult: + anchor_page: int + entries: list[dict[str, Any]] = field(default_factory=list) + toc_pages: list[int] = field(default_factory=list) + hierarchies: list[dict[str, Any]] = field(default_factory=list) + batch_meta: list[dict[str, Any]] = field(default_factory=list) + batch_trace: list[dict[str, Any]] = field(default_factory=list) + error: str | None = None + + +def _render_toc_page_batch( + *, + pdf_path: str, + output_dir: str, + dpi: int, + anchor_page: int, + batch_pages: list[int], + render_lock: Any, + reuse_png_by_page: dict[int, str] | None = None, +) -> list[tuple[int, str]]: + """Render one expand window under the Phase2 global render lock. + + One child process renders the whole window (PDF opened once). Across + anchors, PyMuPDF stays serial to avoid gevent ThreadPool deadlocks. + VLM calls happen outside this lock. + """ + reuse = reuse_png_by_page or {} + planned: dict[int, str] = {} + pages_to_render: list[int] = [] + for page_num in batch_pages: + existing = reuse.get(page_num) + if existing and os.path.isfile(existing): + planned[page_num] = existing else: - items = [] + pages_to_render.append(page_num) + + with render_lock: + if pages_to_render: + result = run_in_child_process( + _render_expand_window_worker, + pdf_path, + pages_to_render, + output_dir, + dpi, + anchor_page, + timeout=120, + ) + for item in result.get("results") or []: + planned[int(item["page"])] = str(item["png_path"]) + + page_pngs: list[tuple[int, str]] = [] + missing: list[int] = [] + for page_num in batch_pages: + png_path = planned.get(page_num) + if not png_path: + missing.append(page_num) + continue + page_pngs.append((page_num, png_path)) + if missing: + raise RuntimeError( + f"TOC expand render missing pages {missing} for anchor {anchor_page}" + ) + return page_pngs - confirmed_pages: set[int] = set() - evidence_by_page: dict[int, TocEvidence] = {} - for item in items: - if not isinstance(item, dict) or "page" not in item: - continue - page = int(item["page"]) - is_toc_start = bool(item.get("is_toc_start")) - if is_toc_start: - confirmed_pages.add(page) - raw_confidence = item.get("confidence") - try: - confidence = ( - float(raw_confidence) - if raw_confidence is not None - else (0.95 if is_toc_start else 0.05) + +def _extract_region_for_anchor( + anchor: TocAnchorPage, + *, + pdf_path: str, + page_count: int, + output_dir: str, + dpi: int, + model: str, + render_lock: Any, +) -> _TocRegionResult: + """Expand + extract for one confirmed TOC start. + + Rounds within a start stay serial (continuation context). Different + starts run concurrently for VLM, but page renders share ``render_lock``. + """ + from app.services.document_agent.tools.vlm_toc_extractor import ( + vlm_extract_toc_batch, + ) + + anchor_page = anchor.page + region_entries: list[dict[str, Any]] = [] + region_toc_pages: list[int] = [] + region_scan_end = anchor_page + batch_meta: list[dict[str, Any]] = [] + batch_trace: list[dict[str, Any]] = [] + + try: + for round_idx in range(MAX_BOUNDARY_ROUNDS): + batch_start = anchor_page + round_idx * BOUNDARY_STEP_PAGES + batch_end = min(batch_start + BOUNDARY_STEP_PAGES - 1, page_count) + if batch_start > page_count: + break + + batch_pages = list(range(batch_start, batch_end + 1)) + logger.info( + "[extract.toc] batch round {}: pages {}-{} for anchor {}", + round_idx, + batch_start, + batch_end, + anchor_page, + ) + + reuse_png_by_page: dict[int, str] = {} + if ( + round_idx == 0 + and anchor.png_path + and os.path.isfile(anchor.png_path) + ): + # Phase1 already rendered the confirmed start page. + reuse_png_by_page[anchor_page] = anchor.png_path + + page_pngs = _render_toc_page_batch( + pdf_path=pdf_path, + output_dir=output_dir, + dpi=dpi, + anchor_page=anchor_page, + batch_pages=batch_pages, + render_lock=render_lock, + reuse_png_by_page=reuse_png_by_page, + ) + + batch_result = vlm_extract_toc_batch( + page_pngs=page_pngs, + model=model, + previous_entries=region_entries if region_entries else None, + ) + batch_meta.append(batch_result.meta) + region_entries.extend(batch_result.all_entries) + region_toc_pages.extend(batch_result.toc_pages) + region_scan_end = batch_end + batch_trace.append( + { + "anchor": anchor_page, + "round": round_idx, + "batch_pages": batch_pages, + "toc_pages": batch_result.toc_pages, + "non_toc_pages": batch_result.non_toc_pages, + "entries_found": len(batch_result.all_entries), + } + ) + + last_page_is_toc = ( + batch_result.page_results + and batch_result.page_results[-1].is_toc + ) + if not last_page_is_toc: + logger.info( + "[extract.toc] boundary found: last page {} is not TOC", + batch_end, ) - except (TypeError, ValueError): - confidence = 0.95 if is_toc_start else 0.05 - evidence_by_page[page] = TocEvidence( - page_index=page, - source="vlm", - confidence=max(0.0, min(1.0, confidence)), - reason=str(item.get("reason") or ""), + break + if batch_end >= page_count: + break + logger.info( + "[extract.toc] last page {} still TOC, expanding window", + batch_end, ) - confirmed = [a for a in anchor_pages if a.page in confirmed_pages] - rejected = [a.page for a in anchor_pages if a.page not in confirmed_pages] - evidence = [ - evidence_by_page.get( - a.page, - TocEvidence( - page_index=a.page, - source="vlm", - confidence=0.05, - reason="VLM response omitted this candidate page", - ), + hierarchies: list[dict[str, Any]] = [] + if region_entries: + hierarchies = vlm_entries_to_toc_hierarchies( + region_entries, + toc_page_nums=region_toc_pages, + scan_end_page=region_scan_end, + page_count=page_count, ) - for a in anchor_pages - ] - logger.info( - "[extract.toc] VLM confirmed {} TOC starts, rejected pages: {}", - len(confirmed), - rejected, + return _TocRegionResult( + anchor_page=anchor_page, + entries=region_entries, + toc_pages=region_toc_pages, + hierarchies=hierarchies, + batch_meta=batch_meta, + batch_trace=batch_trace, ) - return confirmed, False, evidence except Exception as exc: - if budget: - budget.refund("visual", est=est, stage=stage) logger.warning( - "[extract.toc] VLM anchor confirmation failed: {}, " - "falling back to no confirmed anchors (safe degradation)", + "[extract.toc] anchor {} region extract failed: {}", + anchor_page, exc, ) - return [], True, [] + return _TocRegionResult( + anchor_page=anchor_page, + entries=region_entries, + toc_pages=region_toc_pages, + batch_meta=batch_meta, + batch_trace=batch_trace, + error=str(exc), + ) + + +def _extract_regions_for_confirmed_anchors( + confirmed: list[TocAnchorPage], + *, + pdf_path: str, + page_count: int, + output_dir: str, + dpi: int, + model: str, +) -> list[_TocRegionResult]: + """Phase 2: concurrent VLM per start; serial PyMuPDF renders across starts.""" + if not confirmed: + return [] + + from gevent.lock import Semaphore + from gevent.pool import Pool as GeventPool + + pool_size = min(BOUNDARY_STEP_PAGES, len(confirmed)) + render_lock = Semaphore(1) + logger.info( + "[extract.toc] Phase 2 extract: {} confirmed starts, " + "vlm_concurrency={}, render=serial", + len(confirmed), + pool_size, + ) + pool = GeventPool(size=pool_size) + jobs = [ + pool.spawn( + _extract_region_for_anchor, + anchor, + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + dpi=dpi, + model=model, + render_lock=render_lock, + ) + for anchor in confirmed + ] + pool.join() + + by_anchor: dict[int, _TocRegionResult] = {} + for job in jobs: + try: + result = job.get() + except Exception as exc: + logger.warning("[extract.toc] region greenlet failed: {}", exc) + continue + by_anchor[result.anchor_page] = result + + # Preserve document page order when merging regions. + ordered: list[_TocRegionResult] = [] + for anchor in confirmed: + result = by_anchor.get(anchor.page) + if result is None: + ordered.append( + _TocRegionResult( + anchor_page=anchor.page, + error="region greenlet failed", + ) + ) + else: + ordered.append(result) + return ordered # -- Main tool ----------------------------------------------------------------- @@ -261,7 +623,7 @@ def extract_toc_with_boundaries( ) os.makedirs(output_dir, exist_ok=True) - # -- Phase 1: VLM confirm anchors ----------------------------------------- + # -- Phase 1: VLM confirm anchors (batched + concurrent) ------------------- confirmed, confirm_failed, confirm_evidence = _vlm_confirm_anchors( anchors, model, budget=ctx.budget ) @@ -305,14 +667,14 @@ def extract_toc_with_boundaries( debug=debug_info, ) - # -- Phase 2+3 (unified): batch classify + extract --------------------------- - # Instead of separate boundary detection (Phase 2) then per-page extraction - # (Phase 3), we send batches of BOUNDARY_STEP_PAGES images to VLM in one - # call. The VLM classifies each page (TOC vs non-TOC) AND extracts entries - # from TOC pages simultaneously. If the last page in a batch is still TOC, - # we expand the window and use prior entries as continuation context. - from app.services.document_agent.tools.vlm_toc_extractor import ( - vlm_extract_toc_batch, + # -- Phase 2: per-confirmed-start expand + extract (concurrent across starts) + region_results = _extract_regions_for_confirmed_anchors( + confirmed, + pdf_path=ctx.pdf_path, + page_count=page_count, + output_dir=output_dir, + dpi=dpi, + model=model, ) all_entries: list[dict[str, Any]] = [] @@ -321,98 +683,27 @@ def extract_toc_with_boundaries( batch_meta: list[dict[str, Any]] = [] batch_trace: list[dict[str, Any]] = [] - for anchor in confirmed: - anchor_page = anchor.page - region_entries: list[dict[str, Any]] = [] - region_toc_pages: list[int] = [] - region_scan_end = anchor_page - - for round_idx in range(MAX_BOUNDARY_ROUNDS): - batch_start = anchor_page + round_idx * BOUNDARY_STEP_PAGES - batch_end = min( - batch_start + BOUNDARY_STEP_PAGES - 1, page_count - ) - if batch_start > page_count: - break - - batch_pages = list(range(batch_start, batch_end + 1)) - logger.info( - "[extract.toc] batch round {}: pages {}-{} for anchor {}", - round_idx, batch_start, batch_end, anchor_page, - ) - - # Render all pages in this batch - page_pngs: list[tuple[int, str]] = [] - for page_num in batch_pages: - png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") - run_in_child_process( - _render_single_page_worker, - ctx.pdf_path, - page_num, - png_path, - dpi, - timeout=60, - ) - page_pngs.append((page_num, png_path)) - - # Send batch to VLM — classify + extract in one call - batch_result = vlm_extract_toc_batch( - page_pngs=page_pngs, - model=model, - previous_entries=region_entries if region_entries else None, - ) - batch_meta.append(batch_result.meta) - - # Collect results - region_entries.extend(batch_result.all_entries) - region_toc_pages.extend(batch_result.toc_pages) - region_scan_end = batch_end - - batch_trace.append({ - "anchor": anchor_page, - "round": round_idx, - "batch_pages": batch_pages, - "toc_pages": batch_result.toc_pages, - "non_toc_pages": batch_result.non_toc_pages, - "entries_found": len(batch_result.all_entries), - }) - - # Determine if we need to continue expanding - # If the last page in the batch is NOT TOC, boundary found - last_page_is_toc = ( - batch_result.page_results - and batch_result.page_results[-1].is_toc + for region in region_results: + if region.error: + warnings.append( + f"toc_region_failed:anchor={region.anchor_page}:{region.error}" ) - if not last_page_is_toc: - logger.info( - "[extract.toc] boundary found: last page {} is not TOC", - batch_end, - ) - break - - # Last page is still TOC — continue expanding - if batch_end >= page_count: - break - logger.info( - "[extract.toc] last page {} still TOC, expanding window", - batch_end, - ) - - all_entries.extend(region_entries) - all_toc_pages.extend(region_toc_pages) - - if region_entries: - region_hierarchies = vlm_entries_to_toc_hierarchies( - region_entries, - toc_page_nums=region_toc_pages, - scan_end_page=region_scan_end, - page_count=page_count, + logger.warning( + "[extract.toc] anchor {} region failed: {}", + region.anchor_page, + region.error, ) - toc_hierarchies.extend(region_hierarchies) - else: + continue + all_entries.extend(region.entries) + all_toc_pages.extend(region.toc_pages) + batch_meta.extend(region.batch_meta) + batch_trace.extend(region.batch_trace) + if region.hierarchies: + toc_hierarchies.extend(region.hierarchies) + elif not region.entries: logger.warning( "[extract.toc] anchor {} produced no TOC entries", - anchor_page, + region.anchor_page, ) if not all_entries: diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index 1e17a2202..f5865d856 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -10,6 +10,10 @@ from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult +# Debug / package-root artifact name. Keep exported so page_memory debug scripts +# can import it from ``document_agent.persist``. +DOC_PROFILE_FILENAME = "doc_profile.json" + def _artifact_dir(ctx: ToolContext) -> Path: if ctx.output_dir: diff --git a/apps/worker/app/services/document_agent/tools/probe_page_features.py b/apps/worker/app/services/document_agent/tools/probe_page_features.py index be92b853c..2c3de3598 100644 --- a/apps/worker/app/services/document_agent/tools/probe_page_features.py +++ b/apps/worker/app/services/document_agent/tools/probe_page_features.py @@ -203,7 +203,7 @@ def _probe_visual_assets( header_y: float | None, footer_y: float | None, ) -> dict[str, Any]: - """Collect counts + coarse asset bboxes from images / tables / drawings.""" + """Collect counts + has_asset gate from images / tables / drawings.""" image_area = 0.0 bboxes: list[dict[str, Any]] = [] seen_image_rects: set[tuple[float, float, float, float]] = set() @@ -276,7 +276,8 @@ def _probe_visual_assets( "image_count": image_count, "table_count": table_count, "drawings_count": drawings_count, - "asset_bboxes": bboxes or None, + # Geometry is only used to derive the gate; do not persist bboxes. + "has_asset": bool(bboxes), } @@ -364,7 +365,6 @@ def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: height=float(item.get("height") or 0.0), has_asset=False, is_blank_like=bool(item.get("is_blank_like")), - asset_bboxes=None, ) for item in (result.get("features") or []) ] @@ -412,7 +412,7 @@ def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: updated: list[PageFeature] = [] for feature in ctx.blackboard.page_features: visual = by_page.get(feature.page) or {} - has_asset = visual.get("asset_bboxes") is not None + has_asset = bool(visual.get("has_asset")) updated.append( PageFeature( page=feature.page, @@ -427,11 +427,6 @@ def probe_page_assets(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: height=feature.height, has_asset=has_asset, is_blank_like=feature.raw_text_length < 50 and not has_asset, - asset_bboxes=( - list(visual["asset_bboxes"]) - if isinstance(visual.get("asset_bboxes"), list) - else None - ), ) ) ctx.blackboard.page_features = sorted(updated, key=lambda f: f.page) diff --git a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py index 9c6c87c95..cc331b290 100644 --- a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py +++ b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py @@ -8,6 +8,9 @@ from dataclasses import dataclass from typing import Any, cast +# Shared completion budget for TOC VLM calls (confirm batches + extract batches). +TOC_VLM_MAX_TOKENS = 8192 + # --------------------------------------------------------------------------- # Batch-mode prompt: send a window of candidate pages in one VLM call. @@ -196,7 +199,7 @@ def vlm_extract_toc_batch( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, temperature=0.1, - max_tokens=8192, + max_tokens=TOC_VLM_MAX_TOKENS, response_format={"type": "json_object"}, usage_task="document_agent.vlm_toc_batch", ) diff --git a/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py index 06ff1ad06..7684fee73 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py +++ b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py @@ -29,7 +29,7 @@ from shared.services.ai.summary.engine import summarize from shared.services.ai.summary.model import AssetSummary, BodySummary from shared.utils.chunk_refs import build_chunk_ref -from app.services.common.file_utils import path_handle +from app.services.common.file_utils import MAX_ASSET_FILE_NAME_CHARS, path_handle # Each deferred task now carries the engine's typed contract straight through to # the apply step (audit §4.5): assets → AssetSummary, text → BodySummary. The row @@ -280,7 +280,10 @@ def _apply_image_summary_result( image_dir = original_task.image_dir old_img_name = original_task.image_name image_suffix = original_task.image_suffix - safe_title = path_handle(str(img_title), mode="clean_single") + cleaned_title = path_handle(str(img_title), mode="clean_single") + if not isinstance(cleaned_title, str) or not cleaned_title: + return + safe_title = cleaned_title[:MAX_ASSET_FILE_NAME_CHARS] img_num_match = re.match(r"image-(\d+)", str(old_img_name)) img_num = ( img_num_match.group(1) @@ -290,6 +293,8 @@ def _apply_image_summary_result( else "0" ) new_img_name = path_handle(f"image-{img_num}-{safe_title}", mode="clean_single") + if not isinstance(new_img_name, str) or not new_img_name: + return old_path = os.path.join(image_dir, f"{old_img_name}{image_suffix}") new_path = os.path.join(image_dir, f"{new_img_name}{image_suffix}") if old_path == new_path or not os.path.exists(old_path): diff --git a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py index 52305689d..1958bdbf0 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -19,7 +19,7 @@ from loguru import logger from shared.utils.chunk_refs import build_chunk_ref -from app.services.common.file_utils import path_handle +from app.services.common.file_utils import MAX_ASSET_FILE_NAME_CHARS, path_handle @dataclass(frozen=True) @@ -135,6 +135,7 @@ def build_markdown_image_asset( def build_markdown_image_name(*, image_count: int, last_context: str) -> str: image_name_context = path_handle(last_context.strip(), mode="clean_single") if image_name_context: + image_name_context = image_name_context[:MAX_ASSET_FILE_NAME_CHARS] return f"image-{image_count}-{image_name_context}" return f"image-{image_count}" diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py index 0af790288..07fe1c6d0 100644 --- a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -263,6 +263,16 @@ def demote_consecutive_same_level( Note: a candidate the LLM already demoted to -1 naturally breaks the run, because -1 never equals a positive level — it acts like a body separator. + + TODO(heading-demote): Consider an iterative fixpoint pass that re-runs this + demotion while treating already-demoted ids as transparent (skipped for + adjacency only; final level stays -1). That would clean leftover singleton + TOC lines (e.g. ``2.4 foo....11``) after their same-level siblings were + demoted, i.e. pure outline regions with no body placeholders. Do NOT ship + without guarding the false-positive where empty subsections demote first + (``2.4 / 2.4.1 / 2.4.2 / 2.5 / 2.5.1 / [BODY]``) and a later pass then + wrongly merges real sibling parents ``2.4`` and ``2.5`` into one same-level + run. Deferred until we have TOC-page vs empty-subsection regression cases. """ demoted: set[int] = set() run: list[tuple[int, int]] = [] # (row_id, level) for a no-body candidate run diff --git a/apps/worker/app/services/document_parser/tables/table_text_parser.py b/apps/worker/app/services/document_parser/tables/table_text_parser.py index 261bbe653..935f04e6c 100644 --- a/apps/worker/app/services/document_parser/tables/table_text_parser.py +++ b/apps/worker/app/services/document_parser/tables/table_text_parser.py @@ -8,7 +8,7 @@ import pandas as pd from bs4 import BeautifulSoup, Tag -_MAX_TABLE_NAME_CHARS = 80 +from app.services.common.file_utils import MAX_ASSET_FILE_NAME_CHARS def sanitize_table_name_from_header(raw_header_text: str) -> str: @@ -31,8 +31,8 @@ def sanitize_table_name_from_header(raw_header_text: str) -> str: meaningful = [field for field in unique if _is_meaningful_token(field)] result = " ".join(meaningful) - if len(result) > _MAX_TABLE_NAME_CHARS: - result = result[:_MAX_TABLE_NAME_CHARS].rstrip() + if len(result) > MAX_ASSET_FILE_NAME_CHARS: + result = result[:MAX_ASSET_FILE_NAME_CHARS].rstrip() return result diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 43f664029..ce835659d 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -28,9 +28,6 @@ locate_title_compact_strict, resolve_hierarchy_page_ranges, ) -from app.services.document_parser.structure.body_boundary import ( - clean_toc_title, -) from loguru import logger from shared.services.chunks.path_segments import ( append_document_path, @@ -298,7 +295,8 @@ def _range_to_skeleton( ) -> SectionSkeleton: start_page = _clamp_page(item.start_page, page_count) end_page = _clamp_page(item.end_page, page_count) - path_titles = [clean_toc_title(title) or title for title in item.path_titles] + # Keep original TOC titles (incl. numbering) in section_path / HIERARCHY. + path_titles = [str(title).strip() for title in item.path_titles if str(title).strip()] section_path = join_document_path([filename, *path_titles]) parent_path = ( join_document_path([filename, *path_titles[:-1]]) diff --git a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py deleted file mode 100644 index e3db6597e..000000000 --- a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py +++ /dev/null @@ -1,128 +0,0 @@ -from shared.services.retrieval.agentic.evidence.renderer import render_leaf_chunks - - -def test_render_direct_table_chunk_uses_summary_and_asset_url() -> None: - parts: list[str] = [] - render_leaf_chunks( - parts, - [ - { - "chunk_id": "table-1", - "chunk_type": "table", - "content": "
SHOULD NOT LEAK
", - "file_path": "tables/table-企业入驻信息表.html", - "chunk_metadata": { - "summary": "企业入驻信息登记模板", - "keywords": ["企业信息", "入驻管理"], - }, - } - ], - " ", - asset_lookup={"table-1": "http://localhost:4566/table.html?signature=test"}, - ) - - rendered = "\n".join(parts) - assert "[Table: http://localhost:4566/table.html?signature=test]" in rendered - assert "企业入驻信息登记模板" in rendered - assert "企业信息;入驻管理" in rendered - assert "SHOULD NOT LEAK" not in rendered - assert " None: - parts: list[str] = [] - render_leaf_chunks( - parts, - [ - { - "chunk_id": "table-1", - "chunk_type": "table", - "content": "
企业名称
", - "file_path": "tables/table-企业入驻信息表.html", - "source_chunk_path": "企业信息汇总260509 (1).xlsx/企业批量录入", - "chunk_metadata": { - "summary": "table-企业批量录入\n企业入驻信息登记模板", - "keywords": ["企业信息", "入驻管理"], - }, - } - ], - " ", - asset_lookup={"table-1": "http://localhost:4566/table.html?signature=test"}, - ) - - rendered = "\n".join(parts) - assert "[Table: http://localhost:4566/table.html?signature=test]" in rendered - assert "Table path: 企业信息汇总260509 (1).xlsx/企业批量录入" in rendered - assert "Table asset: tables/table-企业入驻信息表.html" in rendered - assert "table-企业批量录入" in rendered - assert "Main columns:" in rendered - assert "企业信息;入驻管理" in rendered - assert "
企业名称
" not in rendered - - -def test_render_connected_table_chunk_includes_asset_url() -> None: - parts: list[str] = [] - render_leaf_chunks( - parts, - [ - { - "chunk_id": "text-1", - "chunk_type": "text", - "content": "见表 [tables/table-1.html]", - "chunk_metadata": { - "connect_to": [ - { - "target": "table-1", - "ref": "[tables/table-1.html]", - } - ] - }, - }, - { - "chunk_id": "table-1", - "chunk_type": "table", - "content": "
SHOULD NOT LEAK
", - "file_path": "tables/table-1.html", - "chunk_metadata": {"summary": "A 表摘要"}, - }, - ], - " ", - asset_lookup={"table-1": "http://localhost:4566/table-1.html?signature=test"}, - ) - - rendered = "\n".join(parts) - assert "[Table: http://localhost:4566/table-1.html?signature=test]" in rendered - assert "A 表摘要" in rendered - assert "SHOULD NOT LEAK" not in rendered - assert " None: - parts: list[str] = [] - render_leaf_chunks( - parts, - [ - { - "chunk_id": "page-node-1", - "chunk_type": "page", - "content": "RAW OCR SHOULD NOT LEAK", - "chunk_metadata": { - "summary": "制度标准总则摘要", - "page_nums": [225, 226], - }, - } - ], - " ", - asset_lookup={ - "page-node-1": "http://localhost:4566/page_pdfs/225-226.pdf?signature=test" - }, - ) - - rendered = "\n".join(parts) - assert "Pages 225-226" in rendered - assert "制度标准总则摘要" in rendered - assert ( - "Page PDF (pages 225-226): " - "http://localhost:4566/page_pdfs/225-226.pdf?signature=test" - ) in rendered - assert "RAW OCR SHOULD NOT LEAK" not in rendered diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 839b61654..ae11c7750 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -58,7 +58,6 @@ def _page_feature(page: int = 1) -> PageFeature: height=72.0, has_asset=False, is_blank_like=False, - asset_bboxes=None, ) @@ -158,6 +157,7 @@ def test_run_lightweight_anatomy_builds_single_shard_without_planner_llm( ) assert list(anatomy_data)[:2] == ["version", "toc_hierarchies"] assert "text_lines_preview" not in anatomy_data["page_features"][0] + assert "asset_bboxes" not in anatomy_data["page_features"][0] trace_data = json.loads((output_dir / "trace.json").read_text(encoding="utf-8")) assert "visual_stages" in trace_data["summary"]["budget"] diff --git a/apps/worker/tests/contract/test_profile_agent_protocol_contract.py b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py new file mode 100644 index 000000000..b16233a5e --- /dev/null +++ b/apps/worker/tests/contract/test_profile_agent_protocol_contract.py @@ -0,0 +1,214 @@ +"""Protocol tests: planner next_action and executor finish ownership.""" + +from __future__ import annotations + +import json +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.executor.react_loop import ( + ReActExecutor, + _parse_decision, +) +from app.services.document_agent.manifest import ( + DocumentProfile, + ReflexionDecision, + ToolContext, +) +from app.services.document_agent.planner.planner import _parse_profile_and_decision +from app.services.document_agent.registry import REGISTRY +from app.services.document_agent import tools as _registered_tools # noqa: F401 +from app.services.document_agent.state import AgentBlackboard + + +def test_planner_verdict_now_falls_through_to_ready_to_shard() -> None: + raw = json.dumps( + { + "is_scanned": True, + "category": "Feasibility Study Report", + "routing_category": "generic", + "category_rationale": "scanned prose", + "language": "zh", + "rationale": "scanned PDF not atlas", + "header_y": None, + "footer_y": None, + "next_action": "verdict_now", + "inspect_pages": [], + "grep_query": "", + } + ) + profile, decision = _parse_profile_and_decision(raw) + assert profile.is_scanned is True + assert decision.action == "tool_call" + assert decision.tool_name == "propose.shard_plan" + assert decision.verdict is None + + +def test_planner_ready_to_shard_proposes_shard_plan() -> None: + raw = json.dumps( + { + "is_scanned": False, + "category": "Report", + "routing_category": "generic", + "language": "en", + "rationale": "enough evidence", + "next_action": "ready_to_shard", + } + ) + _profile, decision = _parse_profile_and_decision(raw) + assert decision.action == "tool_call" + assert decision.tool_name == "propose.shard_plan" + + +def test_planner_inspect_more_maps_to_inspect_pages() -> None: + raw = json.dumps( + { + "is_scanned": False, + "category": "Report", + "routing_category": "generic", + "language": "en", + "rationale": "need more pages", + "next_action": "inspect_more", + "inspect_pages": [3, 8], + } + ) + _profile, decision = _parse_profile_and_decision(raw) + assert decision.tool_name == "inspect.pages" + assert decision.tool_args["pages"] == [3, 8] + + +def test_executor_legacy_verdict_now_without_status_becomes_shard() -> None: + decision = _parse_decision( + json.dumps( + { + "action": "verdict_now", + "rationale": "classification done", + } + ) + ) + assert decision.action == "tool_call" + assert decision.tool_name == "propose.shard_plan" + + +def test_executor_legacy_verdict_now_with_abort_status_uses_verdict_tool() -> None: + decision = _parse_decision( + json.dumps( + { + "action": "verdict_now", + "rationale": "cannot profile", + "verdict": {"status": "abort", "rationale": "cannot profile"}, + } + ) + ) + assert decision.action == "tool_call" + assert decision.tool_name == "verdict" + assert decision.tool_args["status"] == "abort" + + +def _seed_pages(blackboard: AgentBlackboard, page_count: int) -> None: + from app.services.document_agent.manifest import PageFeature, PageLabel + + blackboard.page_count = page_count + blackboard.doc_stats = {"page_count": page_count} + blackboard.page_features = [ + PageFeature( + page=page, + raw_text_length=0, + text_density=0.0, + image_coverage=1.0, + image_count=1, + table_count=0, + drawings_count=0, + orientation="portrait", + width=612.0, + height=792.0, + has_asset=True, + is_blank_like=True, + ) + for page in range(1, page_count + 1) + ] + blackboard.page_labels = [ + PageLabel(page=page, kind="normal", confidence=0.9) + for page in range(1, page_count + 1) + ] + + +def test_executor_initial_ready_to_shard_reaches_success_without_abort() -> None: + blackboard = AgentBlackboard() + _seed_pages(blackboard, 4) + blackboard.document_profile = DocumentProfile( + is_scanned=True, + category="Feasibility Study Report", + routing_category="generic", + rationale="scanned PDF not atlas", + ) + from app.services.document_agent.manifest import TocResult + + blackboard.toc_result = TocResult(method="none", notes="no toc") + ctx = ToolContext( + pdf_path="/tmp/scanned.pdf", + job_id="job-scanned", + blackboard=blackboard, + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, # deterministic executor (no LLM) + ) + initial = ReflexionDecision( + action="tool_call", + rationale="ready", + tool_name="propose.shard_plan", + tool_args={}, + ) + result = ReActExecutor( + ctx, + registry=REGISTRY, + max_rounds=10, + initial_decision=initial, + ).run() + assert result.verdict.status == "success" + assert blackboard.shard_plan is not None + assert len(blackboard.shard_plan.shards) >= 1 + + +def test_executor_empty_initial_tool_falls_through_to_success() -> None: + """Missing tool_name must coerce to propose.shard_plan, not abort.""" + blackboard = AgentBlackboard() + _seed_pages(blackboard, 3) + from app.services.document_agent.manifest import TocResult + + blackboard.toc_result = TocResult(method="none", notes="no toc") + blackboard.document_profile = DocumentProfile( + is_scanned=True, + category="Report", + routing_category="generic", + ) + ctx = ToolContext( + pdf_path="/tmp/scanned.pdf", + job_id="job-legacy", + blackboard=blackboard, + budget=BudgetTracker(plan_budget=50_000, visual_budget=80_000), + trace=None, + settings={}, + ) + initial = ReflexionDecision( + action="tool_call", + rationale="stale empty decision", + tool_name=None, + tool_args={}, + ) + result = ReActExecutor( + ctx, + registry=REGISTRY, + max_rounds=10, + initial_decision=initial, + ).run() + assert result.verdict.status == "success" + assert blackboard.shard_plan is not None + assert len(blackboard.shard_plan.shards) == 1 diff --git a/apps/worker/tests/contract/test_table_embedded_images_contract.py b/apps/worker/tests/contract/test_table_embedded_images_contract.py index 0da70fad8..79cff2e3e 100644 --- a/apps/worker/tests/contract/test_table_embedded_images_contract.py +++ b/apps/worker/tests/contract/test_table_embedded_images_contract.py @@ -32,11 +32,6 @@ from shared.services.chunks.dataframe_chunk_converter import ( # noqa: E402 dataframe_to_chunks, ) -from shared.services.retrieval.agentic.evidence.renderer import ( # noqa: E402 - render_table_chunk_lines, -) - - def _write_jpeg( path: Path, *, @@ -183,28 +178,6 @@ def test_table_embedded_images_are_extracted_rewritten_and_linked( for item in table_connect ) - rendered = render_table_chunk_lines( - { - "chunk_id": table_id, - "chunk_type": "table", - "file_path": by_type["table"]["metadata"].get("file_path"), - "chunk_metadata": by_type["table"]["metadata"], - }, - display_ref="tables/demo.html", - chunk_by_id={ - image_id: { - "chunk_id": image_id, - "chunk_type": "image", - "file_path": by_type["image"]["metadata"].get("file_path"), - "content": by_type["image"]["content"], - } - }, - asset_lookup={image_id: "https://example.com/img.jpg"}, - rendered_ids=set(), - ) - assert any("[Image: https://example.com/img.jpg]" in line for line in rendered) - - def test_table_embedded_images_skip_below_img_min_size(tmp_path: Path) -> None: output_dir = tmp_path / "doc" images_dir = output_dir / "images" diff --git a/apps/worker/tests/contract/test_toc_confirm_batch_contract.py b/apps/worker/tests/contract/test_toc_confirm_batch_contract.py new file mode 100644 index 000000000..f92992519 --- /dev/null +++ b/apps/worker/tests/contract/test_toc_confirm_batch_contract.py @@ -0,0 +1,126 @@ +"""Contract tests for TOC Phase-1 batched VLM anchor confirmation.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from app.services.document_agent.budget import BudgetTracker, StageEnvelope +from app.services.document_agent.manifest import TocAnchorPage +from app.services.document_agent.tools import extract_toc_with_boundaries as toc_tool + + +def _anchors(tmp_path: Path, pages: list[int]) -> list[TocAnchorPage]: + out: list[TocAnchorPage] = [] + for page in pages: + png = tmp_path / f"toc_anchor_page_{page}.png" + png.write_bytes(b"fake-png") + out.append( + TocAnchorPage( + page=page, + png_path=str(png), + source="text_scan", + ) + ) + return out + + +def test_iter_chunks_uses_boundary_step_size() -> None: + items = list(range(12)) + chunks = toc_tool._iter_chunks(items, toc_tool.BOUNDARY_STEP_PAGES) # noqa: SLF001 + assert chunks == [ + [0, 1, 2, 3, 4], + [5, 6, 7, 8, 9], + [10, 11], + ] + + +def test_vlm_confirm_anchors_batches_and_merges_partial_failures( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """One confirm chunk may fail; successful chunks still contribute confirmed pages.""" + anchors = _anchors(tmp_path, [5, 11, 40, 44, 55, 78, 97]) + assert toc_tool.BOUNDARY_STEP_PAGES == 5 + # 7 anchors → 2 chunks: [5..55] and [78, 97] + + call_pages: list[list[int]] = [] + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> tuple[str, dict[str, int]]: + content = kwargs["messages"][0]["content"] + pages = [] + for part in content: + if not isinstance(part, dict) or part.get("type") != "text": + continue + text = str(part.get("text") or "") + if text.startswith("\n--- Page ") and text.endswith(" ---"): + pages.append(int(text[len("\n--- Page ") : -len(" ---")])) + call_pages.append(pages) + if 78 in pages: + raise RuntimeError("simulated truncated JSON") + payload = { + "pages": [ + { + "page": page, + "is_toc_start": page in {5, 11, 40}, + "reason": "ok", + } + for page in pages + ] + } + return json.dumps(payload), {"total_tokens": 100} + + monkeypatch.setattr( + "shared.services.ai.llm_overrides.get_vision_client", + lambda requested_model=None: (_FakeClient(), requested_model or "fake-vlm"), + ) + + budget = BudgetTracker( + plan_budget=50_000, + visual_budget=200_000, + visual_stage_envelopes={ + "toc_confirm": StageEnvelope(min_guarantee=0, cap=None), + }, + ) + confirmed, confirm_failed, evidence = toc_tool._vlm_confirm_anchors( # noqa: SLF001 + anchors, + model="fake-vlm", + budget=budget, + ) + + assert sorted(call_pages[0] + call_pages[1]) == [5, 11, 40, 44, 55, 78, 97] + assert {tuple(pages) for pages in call_pages} == { + (5, 11, 40, 44, 55), + (78, 97), + } + assert confirm_failed is False + assert [a.page for a in confirmed] == [5, 11, 40] + assert {e.page_index for e in evidence} == {5, 11, 40, 44, 55, 78, 97} + + +def test_vlm_confirm_anchors_all_chunks_fail_sets_confirm_failed( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + anchors = _anchors(tmp_path, [5, 11]) + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> tuple[str, dict[str, int]]: + raise RuntimeError("boom") + + monkeypatch.setattr( + "shared.services.ai.llm_overrides.get_vision_client", + lambda requested_model=None: (_FakeClient(), requested_model or "fake-vlm"), + ) + + confirmed, confirm_failed, _evidence = toc_tool._vlm_confirm_anchors( # noqa: SLF001 + anchors, + model="fake-vlm", + budget=None, + ) + assert confirmed == [] + assert confirm_failed is True diff --git a/apps/worker/tests/contract/test_toc_phase2_region_concurrency_contract.py b/apps/worker/tests/contract/test_toc_phase2_region_concurrency_contract.py new file mode 100644 index 000000000..a068b0b88 --- /dev/null +++ b/apps/worker/tests/contract/test_toc_phase2_region_concurrency_contract.py @@ -0,0 +1,276 @@ +"""Contract tests for TOC Phase-2 concurrent VLM with serial batch renders.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import gevent +import pytest + +from app.services.document_agent.manifest import TocAnchorPage +from app.services.document_agent.tools import extract_toc_with_boundaries as toc_tool +from app.services.document_agent.tools.vlm_toc_extractor import ( + BatchPageResult, + BatchTocResult, +) + + +def _anchors(tmp_path: Path, pages: list[int]) -> list[TocAnchorPage]: + out: list[TocAnchorPage] = [] + for page in pages: + png = tmp_path / f"toc_anchor_page_{page}.png" + png.write_bytes(b"anchor-png") + out.append( + TocAnchorPage(page=page, png_path=str(png), source="vlm") + ) + return out + + +def _anchor_from_png(png_path: str) -> int: + name = Path(png_path).name + if name.startswith("toc_anchor_page_"): + return int(name[len("toc_anchor_page_") : -len(".png")]) + # toc_a{anchor}_p{page}.png + stem = name.removesuffix(".png") + anchor_part, _sep, _page_part = stem.partition("_p") + assert anchor_part.startswith("toc_a"), png_path + return int(anchor_part[len("toc_a") :]) + + +def _batch_result( + *, + toc_pages: list[int], + non_toc_pages: list[int], + entries: list[dict[str, Any]] | None = None, +) -> BatchTocResult: + page_results: list[BatchPageResult] = [ + BatchPageResult(page=page, is_toc=True, entries=[]) for page in toc_pages + ] + page_results.extend( + BatchPageResult(page=page, is_toc=False, entries=[]) for page in non_toc_pages + ) + if page_results and page_results[-1].is_toc: + page_results[-1] = BatchPageResult( + page=page_results[-1].page, + is_toc=False, + entries=[], + ) + return BatchTocResult( + page_results=page_results, + toc_pages=list(toc_pages), + non_toc_pages=list(non_toc_pages), + all_entries=list(entries or []), + meta={"ok": True}, + ) + + +def _fake_batch_render(worker_fn: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + """Match ``_render_expand_window_worker`` args: pages list in one spawn.""" + pages = list(args[1]) + output_dir = Path(args[2]) + anchor_page = int(args[4]) + output_dir.mkdir(parents=True, exist_ok=True) + results: list[dict[str, Any]] = [] + for page_num in pages: + png_path = output_dir / f"toc_a{anchor_page}_p{page_num}.png" + png_path.write_bytes(b"png") + results.append({"page": page_num, "png_path": str(png_path)}) + return {"ok": True, "results": results} + + +def test_extract_regions_runs_per_anchor_and_merges_in_page_order( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + confirmed = _anchors(tmp_path, [10, 40, 80]) + calls: list[int] = [] + render_calls: list[list[int]] = [] + + def _tracking_render(worker_fn: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + render_calls.append(list(args[1])) + return _fake_batch_render(worker_fn, *args, **kwargs) + + def _fake_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, + ) -> BatchTocResult: + first_page = page_pngs[0][0] + anchor = _anchor_from_png(page_pngs[0][1]) + calls.append(anchor) + assert previous_entries in (None, []) + return _batch_result( + toc_pages=[first_page], + non_toc_pages=[p for p, _ in page_pngs[1:]], + entries=[ + { + "title": f"Section {anchor}", + "page": first_page, + "level": 1, + } + ], + ) + + monkeypatch.setattr(toc_tool, "run_in_child_process", _tracking_render) + monkeypatch.setattr( + "app.services.document_agent.tools.vlm_toc_extractor.vlm_extract_toc_batch", + _fake_batch, + ) + monkeypatch.setattr( + toc_tool, + "vlm_entries_to_toc_hierarchies", + lambda entries, **kwargs: [ + { + "toc_range": [entries[0]["page"], entries[0]["page"]], + "toc_tree": {entries[0]["title"]: {}}, + } + ], + ) + + regions = toc_tool._extract_regions_for_confirmed_anchors( # noqa: SLF001 + confirmed, + pdf_path="/tmp/doc.pdf", + page_count=100, + output_dir=str(tmp_path / "toc_pages"), + dpi=72, + model="fake-vlm", + ) + + assert [r.anchor_page for r in regions] == [10, 40, 80] + assert all(r.error is None for r in regions) + assert sorted(calls) == [10, 40, 80] + # One spawn per window; start page reused from Phase1 anchor PNG. + assert sorted(render_calls) == [ + [11, 12, 13, 14], + [41, 42, 43, 44], + [81, 82, 83, 84], + ] + + +def test_extract_regions_keeps_success_when_one_anchor_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + confirmed = _anchors(tmp_path, [10, 40]) + + def _fake_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, + ) -> BatchTocResult: + anchor = _anchor_from_png(page_pngs[0][1]) + if anchor == 40: + raise RuntimeError("simulated region VLM failure") + first_page = page_pngs[0][0] + return _batch_result( + toc_pages=[first_page], + non_toc_pages=[p for p, _ in page_pngs[1:]], + entries=[{"title": "Intro", "page": first_page, "level": 1}], + ) + + monkeypatch.setattr(toc_tool, "run_in_child_process", _fake_batch_render) + monkeypatch.setattr( + "app.services.document_agent.tools.vlm_toc_extractor.vlm_extract_toc_batch", + _fake_batch, + ) + monkeypatch.setattr( + toc_tool, + "vlm_entries_to_toc_hierarchies", + lambda entries, **kwargs: [ + { + "toc_range": [entries[0]["page"], entries[0]["page"]], + "toc_tree": {entries[0]["title"]: {}}, + } + ], + ) + + regions = toc_tool._extract_regions_for_confirmed_anchors( # noqa: SLF001 + confirmed, + pdf_path="/tmp/doc.pdf", + page_count=100, + output_dir=str(tmp_path / "toc_pages"), + dpi=72, + model="fake-vlm", + ) + + assert regions[0].anchor_page == 10 + assert regions[0].error is None + assert regions[0].entries + assert regions[1].anchor_page == 40 + assert regions[1].error is not None + assert "simulated region VLM failure" in regions[1].error + + +def test_phase2_serial_batch_render_with_concurrent_vlm( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Batch renders never overlap; VLM calls may overlap across anchors.""" + confirmed = _anchors(tmp_path, [10, 40, 80]) + render_active = 0 + render_max = 0 + vlm_active = 0 + vlm_max = 0 + spawn_count = 0 + + def _fake_render(worker_fn: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + nonlocal render_active, render_max, spawn_count + spawn_count += 1 + render_active += 1 + render_max = max(render_max, render_active) + gevent.sleep(0.02) + result = _fake_batch_render(worker_fn, *args, **kwargs) + render_active -= 1 + return result + + def _fake_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, + ) -> BatchTocResult: + nonlocal vlm_active, vlm_max + vlm_active += 1 + vlm_max = max(vlm_max, vlm_active) + gevent.sleep(0.2) + first_page = page_pngs[0][0] + anchor = _anchor_from_png(page_pngs[0][1]) + vlm_active -= 1 + return _batch_result( + toc_pages=[first_page], + non_toc_pages=[p for p, _ in page_pngs[1:]], + entries=[{"title": f"S{anchor}", "page": first_page, "level": 1}], + ) + + monkeypatch.setattr(toc_tool, "run_in_child_process", _fake_render) + monkeypatch.setattr( + "app.services.document_agent.tools.vlm_toc_extractor.vlm_extract_toc_batch", + _fake_batch, + ) + monkeypatch.setattr( + toc_tool, + "vlm_entries_to_toc_hierarchies", + lambda entries, **kwargs: [ + { + "toc_range": [entries[0]["page"], entries[0]["page"]], + "toc_tree": {entries[0]["title"]: {}}, + } + ], + ) + + regions = toc_tool._extract_regions_for_confirmed_anchors( # noqa: SLF001 + confirmed, + pdf_path="/tmp/doc.pdf", + page_count=100, + output_dir=str(tmp_path / "toc_pages"), + dpi=72, + model="fake-vlm", + ) + + assert all(r.error is None for r in regions) + assert render_max == 1 + assert spawn_count == 3 # one batch spawn per anchor window + assert vlm_max >= 2 diff --git a/packages/shared-python/shared/core/config/__init__.py b/packages/shared-python/shared/core/config/__init__.py index 6889df8a5..5f07b4179 100644 --- a/packages/shared-python/shared/core/config/__init__.py +++ b/packages/shared-python/shared/core/config/__init__.py @@ -19,6 +19,7 @@ from .mineru import MineruConfig from .qstash import QStashConfig from .redis import RedisConfig, RedisConfigManager, RedisPoolManager +from .retrieval import RetrievalConfig from .storage import StorageConfig __all__ = [ @@ -33,6 +34,7 @@ "JobConfig", "AIConfig", "MineruConfig", + "RetrievalConfig", "AppConfig", "app_config", "settings", diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index 487a35808..9cbbdd367 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -44,34 +44,6 @@ class AIConfig(BaseModel): "Same alternates as IMAGE_MODEL (e.g. qwen3-vl-32b-instruct)." ), ) - RETRIEVAL_PLANNER_MODEL: str = Field( - default="", - description="Reasoning-capable model used by the workflow query planner.", - ) - RETRIEVAL_PLANNER_THINKING_BUDGET: int = Field( - default=4000, - description="Token budget for the query planner thinking call.", - ) - RETRIEVAL_DECOMPOSITION_MAX_STEPS: int = Field( - default=5, - description="Maximum number of planned workflow steps.", - ) - RETRIEVAL_WALLET_TOTAL_BUDGET: int = Field( - default=200000, - description="Total workflow token wallet for decomposed retrieval.", - ) - RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET: int = Field( - default=40000, - description="Default token budget issued to each retrieve step.", - ) - RETRIEVAL_WORKFLOW_PARALLEL_MAX: int = Field( - default=3, - description="Maximum concurrent workflow steps in the same DAG batch.", - ) - RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS: float = Field( - default=10.0, - description="Timeout for the optional retrieval workflow planner LLM call.", - ) # Runtime LLM controls. LLM_MOCK_ENABLED: bool = Field( diff --git a/packages/shared-python/shared/core/config/app.py b/packages/shared-python/shared/core/config/app.py index 5ce5f9396..9b03f7889 100644 --- a/packages/shared-python/shared/core/config/app.py +++ b/packages/shared-python/shared/core/config/app.py @@ -13,6 +13,7 @@ from .mineru import MineruConfig from .qstash import QStashConfig from .redis import RedisConfig, RedisConfigManager, RedisPoolManager +from .retrieval import RetrievalConfig from .storage import StorageConfig @@ -27,6 +28,7 @@ class AppConfig( MineruConfig, BillingConfig, JobConfig, + RetrievalConfig, ): """Application configuration — all config components merged.""" diff --git a/packages/shared-python/shared/core/config/retrieval.py b/packages/shared-python/shared/core/config/retrieval.py new file mode 100644 index 000000000..5484c2f67 --- /dev/null +++ b/packages/shared-python/shared/core/config/retrieval.py @@ -0,0 +1,17 @@ +"""Retrieval configuration settings""" + +from pydantic import Field +from pydantic_settings import BaseSettings + + +class RetrievalConfig(BaseSettings): + """Retrieval configuration settings""" + + RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT: int = Field( + default=2000, + ge=1, + description=( + "Maximum rows the Postgres FTS prefilter returns per BM25 channel " + "before Python BM25 reranking. Larger values trade memory for recall." + ), + ) diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py index af8aa3214..75de9c111 100644 --- a/packages/shared-python/shared/services/ai/llm_mock.py +++ b/packages/shared-python/shared/services/ai/llm_mock.py @@ -1,7 +1,5 @@ """Helpers for deterministic mock responses from OpenAI-compatible LLM calls.""" -import json -import re from typing import Any, Dict, List from loguru import logger @@ -20,13 +18,6 @@ def build_mock_chat_completion_response( model_name, task_name, ) - # For agentic tasks that need dynamic content extraction, build the response here. - if task_name == "agentic-planner": - return _build_planner_mock_response(prompt_text) - if task_name == "agentic-navigate": - return _build_navigate_mock_response(prompt_text) - if task_name == "agentic-discovery-select": - return _build_discovery_select_mock_response(prompt_text) return _build_mock_response(task_name) @@ -77,23 +68,6 @@ def _detect_mock_task(prompt_text: str) -> str: """Infer the prompt task so the mock can return a compatible response shape.""" normalized_prompt = prompt_text.lower() - # ── Agentic retrieval prompts (check first — they are structurally distinct) ── - if ( - "you are a retrieval workflow planner" in normalized_prompt - and "concat_final_parts" in normalized_prompt - ): - return "agentic-planner" - if ( - "you are a document navigation agent" in normalized_prompt - and "=== section tree ==" in normalized_prompt - ): - return "agentic-navigate" - if ( - "=== discovery candidates ==" in normalized_prompt - and "\"selections\"" in normalized_prompt - ): - return "agentic-discovery-select" - # ── Document parsing / ingestion prompts ── if ( "generate a concise title" in normalized_prompt @@ -149,135 +123,9 @@ def _detect_mock_task(prompt_text: str) -> str: return "default" -def _extract_first_section_path(prompt_text: str) -> str | None: - """Pull the first path value from a COLLECTOR_PROMPT section tree block. - - The section tree is rendered by section_prompt_projection.format_items_for_llm() - and each item line looks like:: - - ▸ [L1] path="Root" [text=1] ~100 tokens [Leaf] - └ [L2] path="Root / Sub" [text=2] ~200 tokens - - We extract the value inside path="..." from the first matching line - within the === Section Tree === block. - """ - tree_match = re.search( - r"=== Section Tree ===(.*?)=== End Section Tree ===", - prompt_text, - re.DOTALL | re.IGNORECASE, - ) - if not tree_match: - return None - tree_block = tree_match.group(1) - # Match path="..." in the section tree — this is the canonical format - path_match = re.search(r'path="([^"]+)"', tree_block) - if path_match: - return path_match.group(1) - return None - - - -def _extract_user_query(prompt_text: str) -> str: - """Extract the user query line from a planner/navigation prompt. - - The PLANNER_PROMPT and COLLECTOR_PROMPT both contain:: - User query: {query} - """ - match = re.search(r"User query:\s*(.+)", prompt_text) - if match: - return match.group(1).strip() - return "mock query" - - -def _build_planner_mock_response(prompt_text: str) -> str: - """Return a valid single-step QueryPlan JSON using the real query from the prompt.""" - query = _extract_user_query(prompt_text) - response = { - "reasoning_summary": "mock single-step plan", - "steps": [ - { - "id": "s1", - "sub_query": query, - "step_kind": "retrieve", - "depends_on": [], - "output_role": "final_part", - "top_k": 10, - } - ], - "final_strategy": "concat_final_parts", - } - return json.dumps(response) - - -def _build_navigate_mock_response(prompt_text: str) -> str: - """Return a mock COLLECTOR_PROMPT response that COLLECTs the first visible path.""" - path = _extract_first_section_path(prompt_text) - if path: - response = { - "collect": [{"path": path, "confidence": 0.9, "outline": False}], - "action": "STOP", - "drill_into": None, - "tools": [], - "reason": "Mock: collected first available section", - } - else: - # No path found — STOP without collecting (safe fallback) - response = { - "collect": [], - "action": "STOP", - "drill_into": None, - "tools": [], - "reason": "Mock: no section path found in tree", - } - return json.dumps(response) - - -def _extract_first_discovery_path(prompt_text: str) -> str | None: - """Pull the first path value from a DISCOVERY_SELECT_PROMPT candidates block. - - Discovery hints are rendered by selection._project_discovery_hints() as:: - - ▸ path="Findings" - - - We extract the value inside path="..." from the candidates block. - """ - candidates_match = re.search( - r"=== Discovery Candidates ===(.*?)=== End Discovery Candidates ===", - prompt_text, - re.DOTALL | re.IGNORECASE, - ) - if not candidates_match: - return None - block = candidates_match.group(1) - # Match path="..." — same canonical format as section tree - path_match = re.search(r'path="([^"]+)"', block) - if path_match: - return path_match.group(1) - return None - - -def _build_discovery_select_mock_response(prompt_text: str) -> str: - """Return a mock DISCOVERY_SELECT_PROMPT response selecting the first candidate.""" - path = _extract_first_discovery_path(prompt_text) - if path: - response = {"selections": [{"path": path, "confidence": 0.85}]} - else: - response = {"selections": []} - return json.dumps(response) - - def _build_mock_response(task_name: str) -> str: """Return a canned response compatible with the inferred task contract.""" response_by_task: Dict[str, str] = { - # Agentic retrieval — static fallbacks (dynamic responses built elsewhere) - "agentic-planner": ( - '{"reasoning_summary": "mock single-step plan", ' - '"steps": [{"id": "s1", "sub_query": "mock query", ' - '"step_kind": "retrieve", "depends_on": [], ' - '"output_role": "final_part", "top_k": 10}], ' - '"final_strategy": "concat_final_parts"}' - ), # Document parsing / ingestion tasks "fragment-title": "Mock Fragment Title", "detect-toc-range": '{"toc_start": null, "toc_end": null, "confidence": "low"}', diff --git a/packages/shared-python/shared/services/retrieval/agentic/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/__init__.py deleted file mode 100644 index 961ba1dad..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Agentic evidence retrieval orchestration for Knowhere. - -Flow: - Phase 1: Document selection (discovery + KG LLM select) - Phase 2: Per-document iterative navigation (navigate_step) - Phase 3: Render evidence text for downstream agents - -Each navigate_step chooses one observe-act action plus optional collection -side effects. KNOWHERE does not generate final answers; downstream agents -decide whether the evidence is sufficient. -""" diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/budget.py b/packages/shared-python/shared/services/retrieval/agentic/core/budget.py deleted file mode 100644 index ff7c7f2f7..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/core/budget.py +++ /dev/null @@ -1,361 +0,0 @@ -"""Token budget ledger for agentic retrieval runs.""" -from __future__ import annotations - -import asyncio -import copy -from dataclasses import dataclass -from typing import Any, Literal - - -BudgetPoolName = Literal["bootstrap", "planning", "context"] -BudgetStatus = Literal["HEALTHY", "TIGHT", "CRITICAL", "EXHAUSTED"] - - -def status_from_usage(*, remaining: int, used_pct: int) -> BudgetStatus: - if remaining <= 0: - return "EXHAUSTED" - if used_pct >= 80: - return "CRITICAL" - if used_pct >= 50: - return "TIGHT" - return "HEALTHY" - - -def project_budget_snapshot( - snapshot: dict | None, - *, - pool: BudgetPoolName, - additional_tokens: int, -) -> dict | None: - """Return a snapshot projected after an estimated upcoming token cost.""" - if not snapshot: - return snapshot - adjusted = copy.deepcopy(snapshot) - pool_data = adjusted.get(pool) - if not isinstance(pool_data, dict): - return adjusted - - capacity = int(pool_data.get("capacity") or 0) - used = int(pool_data.get("used") or 0) - reserved = int(pool_data.get("reserved") or 0) - projected_used_total = max(0, used + reserved + max(int(additional_tokens), 0)) - used_pct = ( - int(round(projected_used_total * 100 / capacity)) - if capacity > 0 else 100 - ) - remaining = max(capacity - projected_used_total, 0) - pool_data["used_pct"] = used_pct - pool_data["remaining"] = remaining - pool_data["status"] = status_from_usage( - remaining=remaining, - used_pct=used_pct, - ) - return adjusted - - -def budget_status_from_snapshot( - snapshot: dict[str, Any] | None, - *, - pool: BudgetPoolName = "planning", -) -> str: - """Read a pool status from a serialized budget snapshot.""" - if not isinstance(snapshot, dict): - return "UNKNOWN" - pool_data = snapshot.get(pool) - if not isinstance(pool_data, dict): - return "UNKNOWN" - return str(pool_data.get("status") or "UNKNOWN") - - -class BudgetExceeded(Exception): - """Raised when a planned LLM call cannot reserve budget.""" - - def __init__(self, message: str, *, details: dict[str, Any] | None = None) -> None: - super().__init__(message) - self.details = details or {} - - -@dataclass -class BudgetPool: - name: BudgetPoolName - capacity: int - used: int = 0 - reserved: int = 0 - - @property - def remaining(self) -> int: - return max(self.capacity - self.used - self.reserved, 0) - - @property - def used_pct(self) -> int: - if self.capacity <= 0: - return 100 - return int(round((self.used + self.reserved) * 100 / self.capacity)) - - -class BudgetLedger: - """Concurrency-safe ledger with bootstrap/planning/context pools.""" - - def __init__( - self, - *, - total: int, - planning_ratio: float, - bootstrap: int = 2000, - per_doc_min_share: int = 1500, - ) -> None: - total = max(int(total), 1) - bootstrap = max(0, min(int(bootstrap), total)) - remaining = max(total - bootstrap, 0) - planning_ratio = min(max(float(planning_ratio), 0.0), 1.0) - planning_capacity = int(remaining * planning_ratio) - context_capacity = remaining - planning_capacity - - self._lock = asyncio.Lock() - self._pools: dict[BudgetPoolName, BudgetPool] = { - "bootstrap": BudgetPool("bootstrap", bootstrap), - "planning": BudgetPool("planning", planning_capacity), - "context": BudgetPool("context", context_capacity), - } - self._doc_caps: dict[str, int] = {} - self._doc_used: dict[str, int] = {} - self._doc_reserved: dict[str, int] = {} - self._per_doc_min_share = max(int(per_doc_min_share), 0) - self.total_chunks = 0 - self.total_docs = 0 - self.explored_chunks = 0 - self.explored_docs = 0 - self.trimmed_paths: list[dict[str, Any]] = [] - self._overdraft_events: list[dict[str, Any]] = [] - - def remaining(self, pool: BudgetPoolName) -> int: - return self._pools[pool].remaining - - def status(self, pool: BudgetPoolName) -> BudgetStatus: - pool_state = self._pools[pool] - return status_from_usage( - remaining=pool_state.remaining, - used_pct=pool_state.used_pct, - ) - - async def allocate_doc_caps(self, doc_chunks: dict[str, int]) -> None: - """Allocate planning soft caps by document chunk counts.""" - async with self._lock: - self._doc_caps.clear() - self._doc_used.clear() - self._doc_reserved.clear() - if not doc_chunks: - return - - planning_capacity = self._pools["planning"].capacity - total_weight = sum(max(int(count), 1) for count in doc_chunks.values()) - for doc_id, count in doc_chunks.items(): - weight = max(int(count), 1) - weighted = int(planning_capacity * weight / total_weight) - self._doc_caps[doc_id] = min( - planning_capacity, - max(self._per_doc_min_share, weighted), - ) - - async def try_reserve( - self, - pool: BudgetPoolName, - est: int, - doc_id: str | None = None, - *, - priority: Literal["normal", "low"] = "normal", - ) -> bool: - reservation = await self.reserve( - pool, - est, - doc_id=doc_id, - priority=priority, - allow_overdraft=False, - ) - return bool(reservation.get("reserved")) - - async def reserve( - self, - pool: BudgetPoolName, - est: int, - doc_id: str | None = None, - *, - priority: Literal["normal", "low"] = "normal", - allow_overdraft: bool = False, - overdraft_reason: str = "", - ) -> dict[str, Any]: - est = max(int(est), 0) - if est == 0: - return {"reserved": True, "overdraft": False, "failure": None} - - async with self._lock: - pool_state = self._pools[pool] - if priority == "low" and self.status(pool) == "CRITICAL": - failure = self._reserve_failure(pool, est, doc_id, "low_priority_critical") - return {"reserved": False, "overdraft": False, "failure": failure} - - failure_reason = "" - if pool_state.remaining < est: - failure_reason = "pool_remaining_lt_est" - - # Per-doc cap enforcement: prevent one document from consuming - # the entire planning pool. - doc_remaining: int | None = None - if pool == "planning" and doc_id and doc_id in self._doc_caps: - doc_remaining = self._doc_caps[doc_id] - ( - self._doc_used.get(doc_id, 0) - + self._doc_reserved.get(doc_id, 0) - ) - if doc_remaining < est: - failure_reason = "doc_remaining_lt_est" - - if failure_reason: - failure = self._reserve_failure(pool, est, doc_id, failure_reason) - if not allow_overdraft or pool != "planning": - return {"reserved": False, "overdraft": False, "failure": failure} - self._record_overdraft( - pool=pool, - est=est, - doc_id=doc_id, - reason=overdraft_reason or failure_reason, - failure=failure, - ) - - pool_state.reserved += est - if pool == "planning" and doc_id: - self._doc_reserved[doc_id] = self._doc_reserved.get(doc_id, 0) + est - return { - "reserved": True, - "overdraft": bool(failure_reason), - "failure": self._reserve_failure(pool, est, doc_id, failure_reason) - if failure_reason else None, - } - - async def commit( - self, - pool: BudgetPoolName, - *, - actual: int, - est: int, - doc_id: str | None = None, - ) -> None: - actual = max(int(actual), 0) - est = max(int(est), 0) - async with self._lock: - pool_state = self._pools[pool] - reserved_delta = min(est, pool_state.reserved) - pool_state.reserved -= reserved_delta - pool_state.used = max(0, pool_state.used + actual) - - if pool == "planning" and doc_id: - doc_reserved = min(est, self._doc_reserved.get(doc_id, 0)) - if doc_reserved: - self._doc_reserved[doc_id] -= doc_reserved - if self._doc_reserved[doc_id] <= 0: - self._doc_reserved.pop(doc_id, None) - self._doc_used[doc_id] = self._doc_used.get(doc_id, 0) + actual - - async def refund( - self, - pool: BudgetPoolName, - *, - est: int, - doc_id: str | None = None, - ) -> None: - est = max(int(est), 0) - async with self._lock: - pool_state = self._pools[pool] - pool_state.reserved = max(pool_state.reserved - est, 0) - if pool == "planning" and doc_id: - current = self._doc_reserved.get(doc_id, 0) - remaining = max(current - est, 0) - if remaining: - self._doc_reserved[doc_id] = remaining - else: - self._doc_reserved.pop(doc_id, None) - - def mark_explored( - self, - *, - chunks: int = 0, - docs: int = 0, - ) -> None: - self.explored_chunks += max(int(chunks), 0) - self.explored_docs += max(int(docs), 0) - - def snapshot(self) -> dict[str, object]: - snapshot: dict[str, object] = { - name: { - "capacity": pool.capacity, - "used": pool.used, - "reserved": pool.reserved, - "remaining": pool.remaining, - "used_pct": pool.used_pct, - "overdraft": max(pool.used + pool.reserved - pool.capacity, 0), - "status": self.status(name), - } - for name, pool in self._pools.items() - } - if self._overdraft_events: - snapshot["overdraft_events"] = list(self._overdraft_events) - snapshot.update({ - "total_chunks": self.total_chunks, - "total_docs": self.total_docs, - "explored_chunks": min(self.explored_chunks, self.total_chunks) - if self.total_chunks else self.explored_chunks, - "explored_docs": min(self.explored_docs, self.total_docs) - if self.total_docs else self.explored_docs, - "trimmed_paths": list(self.trimmed_paths), - }) - return snapshot - - def _reserve_failure( - self, - pool: BudgetPoolName, - est: int, - doc_id: str | None, - reason: str, - ) -> dict[str, Any]: - pool_state = self._pools[pool] - details: dict[str, Any] = { - "reason": reason, - "pool": pool, - "prompt_est": max(int(est), 0), - "pool_capacity": pool_state.capacity, - "pool_used": pool_state.used, - "pool_reserved": pool_state.reserved, - "pool_remaining": pool_state.remaining, - } - if pool == "planning" and doc_id and doc_id in self._doc_caps: - doc_used = self._doc_used.get(doc_id, 0) - doc_reserved = self._doc_reserved.get(doc_id, 0) - details.update({ - "doc_id": doc_id, - "doc_cap": self._doc_caps[doc_id], - "doc_used": doc_used, - "doc_reserved": doc_reserved, - "doc_remaining": max(self._doc_caps[doc_id] - doc_used - doc_reserved, 0), - }) - return details - - def _record_overdraft( - self, - *, - pool: BudgetPoolName, - est: int, - doc_id: str | None, - reason: str, - failure: dict[str, Any], - ) -> None: - pool_shortfall = max(int(est) - int(failure.get("pool_remaining") or 0), 0) - doc_shortfall = 0 - if "doc_remaining" in failure: - doc_shortfall = max(int(est) - int(failure.get("doc_remaining") or 0), 0) - self._overdraft_events.append({ - "pool": pool, - "doc_id": doc_id, - "prompt_est": max(int(est), 0), - "shortfall": max(pool_shortfall, doc_shortfall), - "reason": reason, - "failure": failure, - }) diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py deleted file mode 100644 index 91d7a9a80..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Runtime setup helpers for agentic retrieval.""" -from __future__ import annotations - - -import json -import os -from typing import Any - -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document, DocumentChunk -from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetPoolName -from shared.services.retrieval.agentic.core.types import AgentRunConfig, AgentState -from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage -from shared.utils.token_estimate import estimate_tokens - - -def build_config_from_env() -> AgentRunConfig: - return AgentRunConfig( - max_nav_steps=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_STEPS", "6")), - latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "30000")), - token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")), - planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")), - bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")), - per_doc_min_share=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE", "1500")), - inventory_aware=os.environ.get("RETRIEVAL_AGENTIC_INVENTORY_AWARE", "true") == "true", - ) - - -async def load_budget_inventory( - db: AsyncSession, - *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], -) -> tuple[int, int, dict[str, int]]: - stmt = ( - select(Document.document_id, func.count(DocumentChunk.id)) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .group_by(Document.document_id) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - - result = await db.execute(stmt) - doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()} - return sum(doc_chunks.values()), len(doc_chunks), doc_chunks - - -class AgentLlmBudget: - def __init__(self, state: AgentState) -> None: - self._state = state - - async def call( - self, - llm_fn: LLMFn, - prompt: Any, - *, - pool: BudgetPoolName, - doc_id: str | None = None, - priority: str = "normal", - allow_overdraft: bool = False, - overdraft_reason: str = "", - ) -> str: - ledger = self._state.ledger - if ledger is None: - return await llm_fn(prompt) - - prompt_text = _stringify_llm_input(prompt) - est = estimate_tokens(prompt_text) - reservation = await ledger.reserve( - pool, - est, - doc_id=doc_id, - priority="low" if priority == "low" else "normal", - allow_overdraft=allow_overdraft, - overdraft_reason=overdraft_reason, - ) - if not reservation.get("reserved"): - raise BudgetExceeded( - f"{pool} budget exhausted", - details=reservation.get("failure") or {}, - ) - - try: - response = await llm_fn(prompt) - except Exception: - await ledger.refund(pool, est=est, doc_id=doc_id) - raise - - usage = current_llm_usage.get() or {} - actual = _extract_actual_tokens(usage, est) - await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id) - return response - - def for_pool(self, llm_fn: LLMFn, *, pool: BudgetPoolName) -> LLMFn: - async def _call(prompt: Any) -> str: - return await self.call(llm_fn, prompt, pool=pool) - - return _call - - def for_document( - self, - llm_fn: LLMFn, - *, - doc_id: str, - step: int = 0, - allow_overdraft: bool = False, - overdraft_reason: str = "", - ) -> LLMFn: - async def _call(prompt: Any) -> str: - return await self.call( - llm_fn, - prompt, - pool="planning", - doc_id=doc_id, - priority="normal", - allow_overdraft=allow_overdraft, - overdraft_reason=overdraft_reason, - ) - - return _call - - def for_discovery( - self, - llm_fn: LLMFn, - *, - doc_id: str, - low_priority: bool, - ) -> LLMFn: - async def _call(prompt: Any) -> str: - return await self.call( - llm_fn, - prompt, - pool="planning", - doc_id=doc_id, - priority="low" if low_priority else "normal", - ) - - return _call - - -def _extract_actual_tokens(usage: dict, est: int) -> int: - """Derive actual token consumption from LLM usage dict. - - Checks ``total_tokens`` first, then sums ``prompt_tokens`` and - ``completion_tokens``. Falls back to the pre-call estimate. - """ - total = usage.get("total_tokens") - if total: - return int(total) - prompt = int(usage.get("prompt_tokens") or 0) - completion = int(usage.get("completion_tokens") or 0) - return (prompt + completion) or est - - -def _stringify_llm_input(prompt: Any) -> str: - if isinstance(prompt, str): - return prompt - try: - return json.dumps(prompt, ensure_ascii=False, default=str) - except Exception: - return str(prompt) diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/trace.py b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py deleted file mode 100644 index 4d7dd2aa9..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/core/trace.py +++ /dev/null @@ -1,249 +0,0 @@ -"""Lightweight trace recorder for agentic retrieval runs. - -Records each run and its steps into retrieval_runs / retrieval_steps tables. -All writes are best-effort — failures are logged but never propagate to the -caller. This ensures trace recording cannot break the retrieval pipeline. -""" -from __future__ import annotations - -import hashlib -import time -from datetime import datetime, timezone -from typing import Any -from uuid import uuid4 - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.agentic.core.types import ( - AgentRunConfig, - DecisionTraceStep, - ToolResult, -) -from shared.services.retrieval.settings import DEFAULT_TOP_K - - -def _now_utc() -> datetime: - return datetime.now(timezone.utc).replace(tzinfo=None) - - -def _query_hash(query: str) -> str: - return hashlib.sha256(query.encode('utf-8')).hexdigest()[:16] - - -class TraceRecorder: - """Records a single agentic retrieval run and its steps. - - Usage:: - - trace = TraceRecorder(db, user_id=..., namespace=..., query=..., config=...) - await trace.create_run() - ... - trace.record_step(action_type, tool_result) - ... - await trace.complete(ranked_rows, router_used) - """ - - def __init__( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - config: AgentRunConfig, - top_k: int = DEFAULT_TOP_K, - chunk_types: set[str] | None = None, - filters: dict[str, Any] | None = None, - parent_run_id: str | None = None, - workflow_step_id: str | None = None, - workflow_plan: dict[str, Any] | None = None, - ) -> None: - self._db = db - self._run_id = f'aret_{uuid4().hex[:12]}' - self._user_id = user_id - self._namespace = namespace - self._query = query - self._config = config - self._top_k = top_k - self._chunk_types = chunk_types - self._filters = filters or {} - self._parent_run_id = parent_run_id - self._workflow_step_id = workflow_step_id - self._workflow_plan = workflow_plan - self._steps: list[dict[str, Any]] = [] - self._start_time = time.monotonic() - self._created = False - - @property - def run_id(self) -> str: - return self._run_id - - async def create_run(self) -> None: - """Insert the retrieval_runs row. Best-effort.""" - try: - from shared.models.database.document import RetrievalRun - - run = RetrievalRun( - run_id=self._run_id, - user_id=self._user_id, - namespace=self._namespace, - query=self._query, - query_hash=_query_hash(self._query), - top_k=self._top_k, - chunk_types=sorted(self._chunk_types) if self._chunk_types else None, - filters=self._filters, - policy_name='llm_policy_v1', - agentic_enabled=True, - cache_hit=False, - result_count=0, - parent_run_id=self._parent_run_id, - workflow_step_id=self._workflow_step_id, - workflow_plan=self._workflow_plan, - latency_ms=0, - created_at=_now_utc(), - ) - self._db.add(run) - await self._db.flush() - self._created = True - except Exception as e: - logger.debug(f'agentic trace: failed to create run {self._run_id}: {e}') - try: - await self._db.rollback() - except Exception as rollback_error: - logger.debug( - f'agentic trace: failed to roll back create failure ' - f'{self._run_id}: {rollback_error}' - ) - - def record_step( - self, - action_type: str, - result: ToolResult, - *, - decision_reason: str = '', - ) -> None: - """Buffer a step record. Flushed on complete().""" - self._steps.append({ - 'step_index': len(self._steps), - 'action_type': action_type, - 'action_input': {'decision_reason': decision_reason} if decision_reason else {}, - 'observation_status': result.status, - 'observation_payload_keys': list(result.payload.keys()) if result.payload else [], - 'latency_ms': result.latency_ms, - 'error': result.error, - 'tokens_used': result.tokens_used, - 'created_at': _now_utc(), - }) - - def record_decision_trace_step(self, step: DecisionTraceStep) -> None: - """Buffer a DB trace row derived from the public decision trace step.""" - result_status = str(step.result.get("status") or "unknown") - self._steps.append({ - "step_index": len(self._steps), - "action_type": f"{step.phase}:{step.agent}:{step.decision.get('action', '')}", - "action_input": { - "public_step_index": step.step_index, - "decision": step.decision, - "scope": step.scope, - "document_id": step.document_id, - "parent_step_index": step.parent_step_index, - }, - "observation_status": result_status, - "observation_payload_keys": list(step.observation.keys()), - "latency_ms": step.elapsed_ms or 0, - "error": step.result.get("error"), - "tokens_used": 0, - "created_at": _now_utc(), - }) - - def record_budget_stop(self, reason: str) -> None: - """Record that the agent loop stopped due to a budget guard.""" - self._steps.append({ - 'step_index': len(self._steps), - 'action_type': f'budget_stop_{reason}', - 'action_input': {}, - 'observation_status': 'budget_stop', - 'observation_payload_keys': [], - 'latency_ms': 0, - 'error': None, - 'tokens_used': 0, - 'created_at': _now_utc(), - }) - - async def complete( - self, - ranked_rows: list[dict[str, Any]], - router_used: str, - budget_snapshot: dict[str, Any] | None = None, - ) -> None: - """Flush all step records and update the run row. Best-effort.""" - if not self._created: - return - - total_latency = int((time.monotonic() - self._start_time) * 1000) - - try: - from shared.models.database.document import RetrievalStep - - for step_data in self._steps: - step = RetrievalStep( - step_id=f'arst_{uuid4().hex[:12]}', - run_id=self._run_id, - step_index=step_data['step_index'], - action_type=step_data['action_type'], - action_input=step_data.get('action_input'), - observation={ - 'status': step_data['observation_status'], - 'payload_keys': step_data['observation_payload_keys'], - 'tokens_used': step_data.get('tokens_used', 0), - }, - latency_ms=step_data['latency_ms'], - token_count=step_data.get('tokens_used', 0), - error=step_data.get('error'), - created_at=step_data['created_at'], - ) - self._db.add(step) - - # Update run row - from sqlalchemy import update - from shared.models.database.document import RetrievalRun - - # Build provenance: which docs contributed to final results - doc_ids_in_result = list({r.get('document_id', '') for r in ranked_rows if r.get('document_id')}) - provenance = { - 'router': router_used, - 'step_count': len(self._steps), - 'final_doc_ids': doc_ids_in_result, - } - if budget_snapshot is not None: - provenance['budget_snapshot'] = budget_snapshot - if self._parent_run_id: - provenance['parent_run_id'] = self._parent_run_id - if self._workflow_step_id: - provenance['workflow_step_id'] = self._workflow_step_id - - stmt = ( - update(RetrievalRun) - .where(RetrievalRun.run_id == self._run_id) - .values( - result_count=len(ranked_rows), - final_doc_ids=doc_ids_in_result, - result_provenance=provenance, - latency_ms=total_latency, - token_count=sum(step.get('tokens_used', 0) for step in self._steps), - completed_at=_now_utc(), - ) - ) - await self._db.execute(stmt) - await self._db.flush() - - except Exception as e: - logger.debug(f'agentic trace: failed to complete run {self._run_id}: {e}') - try: - await self._db.rollback() - except Exception as rollback_error: - logger.debug( - f'agentic trace: failed to roll back completion failure ' - f'{self._run_id}: {rollback_error}' - ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py deleted file mode 100644 index a3bb1293b..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/core/types.py +++ /dev/null @@ -1,364 +0,0 @@ -"""Core data types for agentic retrieval. - -Defines the state machine primitives: configuration, state, actions, -observations, and tool results. All types are plain dataclasses with -no business logic — they are pure data containers. -""" -from __future__ import annotations - -import time -from dataclasses import dataclass, field -from typing import Any, Literal - -from shared.services.retrieval.agentic.core.budget import BudgetLedger - - -@dataclass -class AgentRunConfig: - """Budget and limit configuration for a single agent run.""" - max_nav_steps: int = 6 # max navigation steps per document (no depth limit) - latency_budget_ms: int = 12000 - token_budget_total: int = 40000 - planning_ratio: float = 0.5 - bootstrap_budget: int = 2000 - per_doc_min_share: int = 1500 - inventory_aware: bool = True - - -@dataclass -class ToolResult: - """Unified return type for all agentic tools. - - Every tool returns one of these, regardless of success or failure. - The orchestrator reads ``status`` to decide next action. - """ - status: str - payload: dict[str, Any] = field(default_factory=dict) - latency_ms: int = 0 - error: str | None = None - tokens_used: int = 0 - - -@dataclass -class DocTreeNode: - """Unified navigation result tree for one document. - - Produced by ``navigate_step``. Captures the full - navigation outcome for rendering as a single hierarchy: - - - ``outline_items``: section tree items at this scope level - - ``leaf_content``: hydrated chunk rows keyed by section path - (leaf selections from LLM) - - ``children``: recursive child trees keyed by section path - (non-leaf selections, populated by orchestrator BFS queue) - - ``confidence``: per-selection confidence for trimming - """ - scope_path: str | None = None - - # Outline items at this level (title + summary) - outline_items: list[dict[str, Any]] = field(default_factory=list) - - # Leaf results, keyed by section path: - leaf_content: dict[str, list[dict[str, Any]]] = field(default_factory=dict) - children: dict[str, 'DocTreeNode'] = field(default_factory=dict) - - # Confidence per selection (for trimming) - confidence: dict[str, float] = field(default_factory=dict) - - @staticmethod - def empty(scope_path: str | None = None) -> 'DocTreeNode': - return DocTreeNode(scope_path=scope_path) - - def has_content(self) -> bool: - """Check if this tree has any meaningful content (outline, chunks, or children).""" - if self.outline_items: - return True - if self.leaf_content: - return True - if self.children: - return any(c.has_content() for c in self.children.values()) - return False - - def has_leaf_content(self) -> bool: - """Check if this tree has any actual hydrated chunk content (not just outline). - - Unlike ``has_content()`` which returns True for outline-only trees, - this method only returns True when real text/table/image chunks have - been hydrated into leaf_content. - """ - if self.leaf_content: - return True - return any(c.has_leaf_content() for c in self.children.values()) - - def flatten_chunk_rows(self) -> list[dict[str, Any]]: - """Recursively collect all hydrated chunk rows (document order).""" - rows: list[dict[str, Any]] = [] - for chunks in self.leaf_content.values(): - rows.extend(chunks) - for child in self.children.values(): - rows.extend(child.flatten_chunk_rows()) - return rows - - def add_leaf_chunks(self, path: str, chunks: list[dict[str, Any]]) -> None: - """Merge chunks into a leaf path, deduplicating by (chunk_id, path).""" - if not path or not chunks: - return - existing = self.leaf_content.setdefault(path, []) - seen: set[tuple[str, str]] = { - (str(row.get('chunk_id') or ''), path) - for row in existing - if row.get('chunk_id') - } - for chunk in chunks: - chunk_id = str(chunk.get('chunk_id') or '') - key = (chunk_id, path) - if chunk_id and key in seen: - continue - if chunk_id: - seen.add(key) - existing.append(chunk) - - def reparent_leaf_content(self) -> None: - """Move descendant leaf paths into matching child nodes. - - Only moves true descendants (prefix match). Content whose path - exactly equals a child key stays here — the renderer handles the - case where a path is both a child and a leaf (section with own - content *and* sub-sections). - """ - for child_path, child in list(self.children.items()): - for leaf_path in list(self.leaf_content.keys()): - if leaf_path.startswith(child_path + ' / '): - child.add_leaf_chunks(leaf_path, self.leaf_content.pop(leaf_path)) - child.reparent_leaf_content() - - def collect_referenced_ids(self, *, document_name: str = '') -> list[dict[str, Any]]: - """Extract minimal chunk references from all hydrated leaves. - - Returns deduplicated list of {chunk_id, document_id, chunk_type, - section_path, file_path, job_id, score} for hit stats and frontend - display. ``score`` carries the real retrieval score (discovery RRF or - navigation confidence) so callers can surface it in results. - """ - refs: list[dict[str, Any]] = [] - seen: set[str] = set() - for row in self.flatten_chunk_rows(): - cid = row.get('chunk_id', '') - if cid and cid not in seen: - seen.add(cid) - section_path = row.get('section_path', '') - if section_path == 'Root' and document_name: - section_path = document_name - refs.append({ - 'chunk_id': cid, - 'document_id': row.get('document_id', ''), - 'chunk_type': row.get('chunk_type', ''), - 'section_path': section_path, - 'file_path': row.get('file_path', ''), - 'job_id': row.get('job_id', ''), - 'score': row.get('score'), # None if no real score available - }) - return refs - - def merge(self, other: 'DocTreeNode') -> None: - """Additive merge for navigation results. - - Merges outline items, leaf content, children, and confidence from - ``other`` into this node. Existing data is preserved; new data is - added. For confidence values, the higher value wins. - """ - existing_paths = {item['path'] for item in self.outline_items} - for item in other.outline_items: - if item.get('path', '') not in existing_paths: - self.outline_items.append(item) - for path, chunks in other.leaf_content.items(): - self.add_leaf_chunks(path, chunks) - for path, child in other.children.items(): - if path in self.children: - self.children[path].merge(child) - else: - self.children[path] = child - for path, conf in other.confidence.items(): - self.confidence[path] = max(self.confidence.get(path, 0), conf) - self.reparent_leaf_content() - - -NavAction = Literal[ - "EXPAND", - "BACK", - "FINISH", - "SEARCH_IMAGES", - "SEARCH_TABLES", - "ERROR", -] - - -@dataclass -class DecisionTraceStep: - """Uniform observe-act-result trace entry exposed to downstream agents.""" - - step_index: int - agent: str - phase: str - observation: dict[str, Any] - decision: dict[str, Any] - result: dict[str, Any] - parent_step_index: int | None = None - document_id: str | None = None - document: str | None = None - scope: str | None = None - budget: dict[str, Any] | None = None - elapsed_ms: int | None = None - - def to_dict(self) -> dict[str, Any]: - data: dict[str, Any] = { - "step_index": self.step_index, - "agent": self.agent, - "parent_step_index": self.parent_step_index, - "phase": self.phase, - "document_id": self.document_id, - "document": self.document, - "scope": self.scope, - "observation": self.observation, - "decision": self.decision, - "result": self.result, - } - if self.budget is not None: - data["budget"] = self.budget - if self.elapsed_ms is not None: - data["elapsed_ms"] = self.elapsed_ms - return data - - -@dataclass -class NavigateStepResult: - """Return type for navigate_step — Collector Agent model. - - Each step returns: - - ``collect``: paths to add to the evidence collection (full hydration) - - ``drill``: paths to explore deeper in subsequent steps - - ``action``: one explicit action — EXPAND/BACK/FINISH/SEARCH_*/ERROR - - ``node``: outline tree node for rendering context - - ``reason``: LLM reasoning for trace - - ``error_reason``: set when action is ERROR — distinguishes system - errors from intentional FINISH so callers can decide retry vs skip. - - ``search_assets_params``: parameters for SEARCH_IMAGES/SEARCH_TABLES - - ``observation``: what the navigator saw before choosing the action - - ``result_status`` / ``result_note``: executor-visible action validation - """ - action: NavAction = "FINISH" - collect: list[dict[str, Any]] = field(default_factory=list) - drill: list[dict[str, Any]] = field(default_factory=list) - back_to: str | None = None # BACK target ancestor path (None = root) - tools: list[str] = field(default_factory=list) - node: DocTreeNode = field(default_factory=DocTreeNode) - reason: str = "" - error_reason: str | None = None - search_assets_params: dict[str, Any] | None = None - observation: dict[str, Any] = field(default_factory=dict) - result_status: str = "ok" - result_note: str | None = None - - @property - def drill_into(self) -> str | None: - """Single drill target path, or None.""" - return self.drill[0]["path"] if self.drill else None - - @property - def is_terminal(self) -> bool: - """True only for explicit terminal actions.""" - return self.action in ("FINISH", "ERROR") - - @staticmethod - def stop(scope_path: str | None = None, *, reason: str = "") -> 'NavigateStepResult': - return NavigateStepResult( - action="FINISH", - node=DocTreeNode.empty(scope_path), - reason=reason, - ) - - @staticmethod - def error(scope_path: str | None = None, *, reason: str = "") -> 'NavigateStepResult': - """Return an ERROR result distinguishable from intentional FINISH.""" - return NavigateStepResult( - action="ERROR", - node=DocTreeNode.empty(scope_path), - reason=f"navigation_error: {reason[:200]}" if reason else "navigation_error", - error_reason=reason[:500] if reason else "unknown_error", - result_status="error", - result_note=reason[:500] if reason else "unknown_error", - ) - - -@dataclass -class CandidateDoc: - """A document selected by kg_document_select.""" - document_id: str - source_file_name: str = '' - confidence: float = 0.0 - reason: str = '' - source: str = '' # 'kg_llm_select' | 'grep' | 'edge_expand' | 'discovery_hint' - - -@dataclass -class AgenticResult: - """Output of agentic retrieval. - - - ``evidence_text``: complete hierarchical context for LLM answering - (rendered doc tree with outline + leaf content + inline tables) - - ``answer_text``: deprecated; always empty because KNOWHERE returns - evidence only and downstream agents synthesize answers. - - ``referenced_chunks``: minimal chunk references for hit stats - and frontend display (chunk_id, document_id, chunk_type, etc.) - - ``router_used``: routing path identifier - - ``budget_snapshot``: final budget ledger state at run completion - - ``stop_reason``: why the run terminated (evidence_only / - budget / latency / max_steps / error / llm_stop) - - ``failure_reason``: fatal retrieval failure reason, if any. - - ``decision_trace``: per-step navigation decisions with reasons, - exposed to downstream agents for stop/retry/modify-query decisions. - """ - evidence_text: str - answer_text: str = '' - referenced_chunks: list[dict[str, Any]] = field(default_factory=list) - router_used: str = 'agentic_discovery_only' - budget_snapshot: dict[str, Any] | None = None - stop_reason: str = '' - failure_reason: str = '' - decision_trace: list[dict[str, Any]] = field(default_factory=list) - - -@dataclass -class AgentState: - """Mutable state carried through the 2-phase orchestrator. - - Phase 1: Document selection (discovery + KG) - Phase 2: Per-document navigation (navigate_step per doc) - Phase 3: Assembly + final verdict - """ - # Timing - start_time: float = field(default_factory=time.monotonic) - step_count: int = 0 - - # Phase 1: Discovery - discovery_top_doc_ids: list[str] = field(default_factory=list) - - # Phase 1: KG document selection - selected_docs: list[CandidateDoc] = field(default_factory=list) - doc_id_to_name: dict[str, str] = field(default_factory=dict) - doc_job_map: dict[str, str] = field(default_factory=dict) - - # Phase 2: Per-document navigation results - doc_trees: dict[str, DocTreeNode] = field(default_factory=dict) # doc_id → DocTreeNode - - ever_explored_doc_ids: set[str] = field(default_factory=set) - - # Token budget + KG inventory - ledger: BudgetLedger | None = None - kg_total_chunks: int = 0 - kg_total_docs: int = 0 - explored_chunks: int = 0 - - @property - def elapsed_ms(self) -> int: - return int((time.monotonic() - self.start_time) * 1000) diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py deleted file mode 100644 index b3fc54a08..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py +++ /dev/null @@ -1,191 +0,0 @@ -"""Discovery and document selection phase for agentic retrieval.""" -from __future__ import annotations - -from typing import Any - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.agentic import tools -from shared.services.retrieval.agentic.core.budget import BudgetExceeded -from shared.services.retrieval.agentic.core.trace import TraceRecorder -from shared.services.retrieval.agentic.core.types import AgentState, CandidateDoc, ToolResult -from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.search.lexical_text import normalize_section_path - - -async def run_initial_discovery( - db: AsyncSession, - *, - state: AgentState, - trace: TraceRecorder, - trace_enabled: bool, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - chunk_types: set[str] | None, - signal_paths: list[str] | None, - filter_mode: str, - channels: list[str] | None, - channel_weights: dict[str, float] | None, - internal_recall_k: int | None, - bootstrap_llm_fn: LLMFn | None, -) -> list[dict[str, Any]]: - discovery_kwargs: dict[str, Any] = { - "user_id": user_id, - "namespace": namespace, - "query": query, - "top_k": top_k, - "exclude_document_ids": exclude_document_ids, - "exclude_sections": exclude_sections, - "chunk_types": chunk_types, - "signal_paths": signal_paths, - "filter_mode": filter_mode, - "channels": channels, - "channel_weights": channel_weights, - "internal_recall_k": internal_recall_k, - } - - logger.info(" agentic: Phase 1 — discovery + document selection") - discovery_result = await tools.bottom_discovery(db, **discovery_kwargs) - state.step_count += 1 - discovery_rows = ( - discovery_result.payload.get("fused_rows", []) - if discovery_result.status != "error" - else [] - ) - state.discovery_top_doc_ids = ( - discovery_result.payload.get("top_doc_ids", []) - if discovery_result.status != "error" - else [] - ) - - if trace_enabled: - trace.record_step( - "bottom_discovery", - discovery_result, - decision_reason="phase_1_mandatory", - ) - - logger.info( - f" agentic step {state.step_count}: bottom_discovery " - f"status={discovery_result.status} latency={discovery_result.latency_ms}ms" - ) - - # Build per-document discovery signals for KG soft-prompting - discovery_signals = build_discovery_signals(discovery_rows) - - if bootstrap_llm_fn is not None: - await _select_documents( - db, - state=state, - trace=trace, - trace_enabled=trace_enabled, - user_id=user_id, - namespace=namespace, - query=query, - exclude_document_ids=exclude_document_ids, - bootstrap_llm_fn=bootstrap_llm_fn, - discovery_signals=discovery_signals, - ) - - return discovery_rows - - -def build_discovery_signals( - discovery_rows: list[dict[str, Any]], -) -> dict[str, list[str]]: - """Build per-document discovery signals from bottom discovery results. - - Returns a mapping of ``{doc_id: [path1, path2, ...]}`` for documents - where keyword/semantic search found potentially relevant section paths. - These signals are injected as soft hints into the KG document selection - prompt, allowing the LLM to make an informed decision rather than - force-injecting documents. - """ - signals: dict[str, list[str]] = {} - seen: dict[str, set[str]] = {} - for row in discovery_rows: - doc_id = row.get("document_id", "") - section_path = normalize_section_path( - str(row.get("section_path", "") or "").strip() - ) - if not doc_id or not section_path or section_path == "Root": - continue - if doc_id not in seen: - seen[doc_id] = set() - signals[doc_id] = [] - if section_path not in seen[doc_id]: - seen[doc_id].add(section_path) - signals[doc_id].append(section_path) - return signals - - -async def _select_documents( - db: AsyncSession, - *, - state: AgentState, - trace: TraceRecorder, - trace_enabled: bool, - user_id: str, - namespace: str, - query: str, - exclude_document_ids: list[str], - bootstrap_llm_fn: LLMFn, - discovery_signals: dict[str, list[str]] | None = None, -) -> None: - try: - kg_result = await tools.kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=bootstrap_llm_fn, - exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - discovery_signals=discovery_signals, - ) - except BudgetExceeded: - logger.info(" agentic: bootstrap budget exhausted during document selection") - if trace_enabled: - trace.record_budget_stop("bootstrap_exhausted") - kg_result = ToolResult( - status="no_confident_doc", - payload={"reason": "bootstrap budget exhausted"}, - ) - state.step_count += 1 - - if trace_enabled: - trace.record_step( - "kg_document_select", - kg_result, - decision_reason="phase_1_doc_selection", - ) - - _append_selected_docs(state, kg_result) - - logger.info( - f" agentic step {state.step_count}: kg_document_select " - f"status={kg_result.status} docs={len(state.selected_docs)} " - f"latency={kg_result.latency_ms}ms" - ) - - -def _append_selected_docs(state: AgentState, kg_result: ToolResult) -> None: - if kg_result.status != "selected_docs": - return - for doc_data in kg_result.payload.get("candidate_docs", []): - state.selected_docs.append( - CandidateDoc( - document_id=doc_data.get("document_id", ""), - source_file_name=doc_data.get("source_file_name", ""), - confidence=doc_data.get("confidence", 0.0), - reason=doc_data.get("reason", ""), - source=doc_data.get("source", ""), - ) - ) - state.doc_id_to_name.update(kg_result.payload.get("doc_id_to_name", {})) - state.doc_job_map.update(kg_result.payload.get("doc_job_map", {})) diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py deleted file mode 100644 index ae62618b7..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py +++ /dev/null @@ -1,329 +0,0 @@ -"""Agentic retrieval discovery tools. - -This Module owns phase-1 retrieval: lexical bottom discovery and document -selection from the document-level knowledge map. The public tool adapter stays -in ``tools.py`` so orchestrator call sites keep a stable interface. -""" -from __future__ import annotations - -import time -from typing import Any - -from loguru import logger -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document -from shared.services.retrieval.agentic.core.budget import BudgetExceeded -from shared.services.retrieval.agentic.navigation.knowledge_map import build_knowledge_map_overview -from shared.services.retrieval.agentic.prompts import ( - FILE_SELECT_PROMPT, - format_budget_block, - parse_json_array, -) -from shared.services.retrieval.agentic.core.types import ToolResult -from shared.services.retrieval.search.channels import content_channel, path_channel, term_channel -from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.search.scoring import ( - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, -) -from shared.services.retrieval.settings import ( - CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER, -) - - -async def bottom_discovery( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], - chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, - internal_recall_k: int | None = None, - **_kwargs: Any, -) -> ToolResult: - """Run 3-channel BM25 discovery plus RRF fusion.""" - t0 = time.monotonic() - try: - allowed_chunk_types = chunk_types - effective_recall_k = ( - internal_recall_k - if internal_recall_k is not None - else top_k * INTERNAL_RECALL_K_MULTIPLIER - ) - active_channels = set(channels) if channels else {"path", "content", "term"} - - path_rows: list[dict[str, Any]] = [] - content_rows: list[dict[str, Any]] = [] - term_rows: list[dict[str, Any]] = [] - - if "path" in active_channels: - path_rows = await path_channel( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - ) - - if "content" in active_channels: - content_rows = await content_channel( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - ) - - if "term" in active_channels: - term_rows = await term_channel( - db, - user_id=user_id, - namespace=namespace, - query=query, - top_k=effective_recall_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - allowed_chunk_types=allowed_chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - ) - - default_weights = { - "path": CHANNEL_WEIGHT_PATH, - "content": CHANNEL_WEIGHT_CONTENT, - "term": CHANNEL_WEIGHT_TERM, - } - effective_weights = {**default_weights, **(channel_weights or {})} - - channel_lists: list[list[dict[str, Any]]] = [] - weight_list: list[float] = [] - if path_rows: - channel_lists.append(path_rows) - weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH)) - if content_rows: - channel_lists.append(content_rows) - weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT)) - if term_rows: - channel_lists.append(term_rows) - weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM)) - - fused_rows = ( - merge_channels_rrf(channel_lists, weight_list, effective_recall_k) - if channel_lists - else [] - ) - fused_rows = merge_same_section_rows(fused_rows) - - if fused_rows: - normalize_row_scores( - fused_rows, - source_field="score", - target_field="discovery_score", - default=0.5, - ) - - doc_id_counts: dict[str, int] = {} - for row in fused_rows: - did = row.get("document_id", "") - if did: - doc_id_counts[did] = doc_id_counts.get(did, 0) + 1 - top_doc_ids = sorted( - doc_id_counts, - key=lambda document_id: doc_id_counts[document_id], - reverse=True, - )[:5] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" agentic.bottom_discovery: {len(fused_rows)} fused rows, " - f"top_doc_ids={top_doc_ids}, {latency}ms" - ) - return ToolResult( - status="discovery_done", - payload={ - "fused_rows": fused_rows, - "top_doc_ids": top_doc_ids, - "channel_counts": { - "path": len(path_rows), - "content": len(content_rows), - "term": len(term_rows), - }, - }, - latency_ms=latency, - ) - except Exception as exc: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f" agentic.bottom_discovery failed: {exc}") - return ToolResult(status="error", error=str(exc), latency_ms=latency) - - -async def kg_document_select( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - llm_fn: LLMFn | None, - exclude_document_ids: list[str], - discovery_signals: dict[str, list[str]] | None = None, - **_kwargs: Any, -) -> ToolResult: - """Select candidate documents from document-level KG.""" - t0 = time.monotonic() - try: - overview_text, doc_id_to_name = await build_knowledge_map_overview( - db, - user_id=user_id, - namespace=namespace, - ) - if overview_text == "(empty)": - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status="no_confident_doc", - payload={"reason": "no active documents in namespace"}, - latency_ms=latency, - ) - - if llm_fn is None: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status="no_confident_doc", - payload={"reason": "LLM not available"}, - latency_ms=latency, - ) - - # Inject discovery signals as soft hints into the overview - if discovery_signals: - overview_text = _inject_discovery_signals( - overview_text, discovery_signals, - ) - - file_prompt = FILE_SELECT_PROMPT.format( - overview=overview_text, - query=query, - budget_block=format_budget_block(_kwargs.get("budget_snapshot")), - ) - file_response = await llm_fn(file_prompt) - selected_ids = parse_json_array(file_response) - - exclude_set = set(exclude_document_ids) - valid_ids = [ - document_id - for document_id in selected_ids - if document_id in doc_id_to_name and document_id not in exclude_set - ] - - if not valid_ids: - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" agentic.kg_document_select: LLM returned no valid docs, {latency}ms" - ) - return ToolResult( - status="no_confident_doc", - payload={ - "reason": "LLM returned no valid document IDs", - "raw_ids": selected_ids, - }, - latency_ms=latency, - ) - - doc_job_map: dict[str, str] = {} - doc_stmt = ( - select(Document.document_id, Document.current_job_result_id) - .where(Document.document_id.in_(valid_ids)) - ) - doc_result = await db.execute(doc_stmt) - for document_id, job_result_id in doc_result.all(): - if job_result_id: - doc_job_map[document_id] = job_result_id - - candidate_docs = [ - { - "document_id": document_id, - "source_file_name": doc_id_to_name.get(document_id, ""), - "confidence": 1.0, - "reason": "LLM selected from KG overview", - "source": "kg_llm_select", - } - for document_id in valid_ids - ] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms" - ) - return ToolResult( - status="selected_docs", - payload={ - "candidate_docs": candidate_docs, - "doc_id_to_name": doc_id_to_name, - "doc_job_map": doc_job_map, - }, - latency_ms=latency, - ) - except BudgetExceeded: - raise - except Exception as exc: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f" agentic.kg_document_select failed: {exc}") - return ToolResult(status="error", error=str(exc), latency_ms=latency) - - -def _inject_discovery_signals( - overview_text: str, - signals: dict[str, list[str]], - *, - max_paths_per_doc: int = 5, -) -> str: - """Inject discovery hint lines into the document overview text. - - For each document that has discovery signals, append hint lines - directly after the document's overview entry. The LLM sees these - as advisory information — it is free to select or ignore the document. - """ - if not signals: - return overview_text - - lines = overview_text.split("\n") - result: list[str] = [] - for line in lines: - result.append(line) - # Match overview lines: "- [doc_xxx] filename chunks=..." - if not line.startswith("- ["): - continue - bracket_end = line.find("]") - if bracket_end < 0: - continue - doc_id = line[3:bracket_end] - paths = signals.get(doc_id) - if not paths: - continue - display_paths = paths[:max_paths_per_doc] - hints_line = ", ".join(f'"{p}"' for p in display_paths) - if len(paths) > max_paths_per_doc: - hints_line += f" (+{len(paths) - max_paths_per_doc} more)" - result.append(f" 🔍 Discovery hints: {hints_line}") - return "\n".join(result) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py deleted file mode 100644 index 80a94954e..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py +++ /dev/null @@ -1,301 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from loguru import logger -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import RetrievalHitStat -from shared.services.retrieval.agentic.core.budget import BudgetLedger -from shared.services.retrieval.agentic.core.types import DocTreeNode -from shared.services.retrieval.hydration.assets import ( - AssetUrlValue, - build_retrieval_asset_url_map, -) -from shared.services.retrieval.stats.service import compute_importance_score -from shared.utils.token_estimate import estimate_tokens - - -def _collect_chunks_by_type( - node: DocTreeNode, - chunk_types: set[str], -) -> list[dict[str, Any]]: - collected_chunks: list[dict[str, Any]] = [] - for chunks in node.leaf_content.values(): - for chunk in chunks: - chunk_type = ( - chunk.get("chunk_type") or chunk.get("type") or "" - ).strip().lower() - if chunk_type in chunk_types: - collected_chunks.append(chunk) - for child in node.children.values(): - collected_chunks.extend(_collect_chunks_by_type(child, chunk_types)) - return collected_chunks - - -def collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: - return _collect_chunks_by_type(node, {"image", "table", "page"}) - - -def collect_media_chunks_all( - doc_trees: dict[str, DocTreeNode], -) -> list[dict[str, Any]]: - media: list[dict[str, Any]] = [] - for tree in doc_trees.values(): - media.extend(collect_media_chunks(tree)) - return media - - -async def build_asset_url_map( - media_chunks: list[dict[str, Any]], -) -> dict[str, AssetUrlValue]: - return await build_retrieval_asset_url_map( - media_chunks, - log_context="agentic evidence", - ) - - -def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: - paths = set(node.leaf_content.keys()) - for child in node.children.values(): - paths.update(_collect_all_leaf_paths(child)) - return paths - - -def _collect_visible_paths(node: DocTreeNode) -> set[str]: - paths = {item["path"] for item in node.outline_items if item.get("path")} - for child in node.children.values(): - paths.update(_collect_visible_paths(child)) - return paths - - -def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None: - parts = path.split(" / ") - for i in range(len(parts) - 1, 0, -1): - ancestor = " / ".join(parts[:i]) - if ancestor in target_paths: - return ancestor - return None - - -def reconcile_deferred_assets( - tree: DocTreeNode, - pending_assets: list[dict], -) -> None: - final_paths = _collect_all_leaf_paths(tree) - visible_paths = _collect_visible_paths(tree) - all_target_paths = final_paths | visible_paths - - if not all_target_paths: - return - - existing_ids = { - str(row.get("chunk_id") or "") - for row in tree.flatten_chunk_rows() - if row.get("chunk_id") - } - - placed = 0 - ancestor_placed = 0 - for asset in pending_assets: - chunk_id = str(asset.get("chunk_id") or "") - if chunk_id and chunk_id in existing_ids: - continue - - owner_path = asset.get("owner_section_path") or asset.get("section_path") - if not owner_path: - continue - - target_path = owner_path if owner_path in all_target_paths else None - if target_path is None: - target_path = _find_closest_ancestor(owner_path, all_target_paths) - if target_path: - ancestor_placed += 1 - - if target_path is None: - continue - - tree.add_leaf_chunks(target_path, [asset]) - if chunk_id: - existing_ids.add(chunk_id) - placed += 1 - - if placed: - tree.reparent_leaf_content() - logger.info( - f" deferred asset reconcile: {placed}/{len(pending_assets)} " - f"assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths " - f"(ancestor_fallback={ancestor_placed})" - ) - - -async def render_evidence( - db: AsyncSession, - doc_trees: dict[str, DocTreeNode], - doc_id_to_name: dict[str, str], -) -> str: - del db - - from shared.services.retrieval.agentic.evidence.renderer import render_unified_doc_tree - - asset_url_map = await build_asset_url_map(collect_media_chunks_all(doc_trees)) - - evidence_parts: list[str] = [] - for doc_id, doc_tree in doc_trees.items(): - # Only render if there is actual hydrated evidence (chunks collected - # via COLLECT or discovery). Outline-only trees (e.g. navigation - # FINISH with empty collect must not leak into evidence_text. - if not doc_tree.has_leaf_content(): - continue - doc_name = doc_id_to_name.get(doc_id, doc_id) - rendered = render_unified_doc_tree( - doc_tree, - doc_name, - asset_lookup=asset_url_map, - ) - if rendered.strip(): - evidence_parts.append(rendered) - - return "\n\n".join(evidence_parts) if evidence_parts else "" - - -def _iter_leaf_content(node: DocTreeNode): - for path, chunks in node.leaf_content.items(): - yield path, chunks - for child in node.children.values(): - yield from _iter_leaf_content(child) - - -def _collect_confidences(node: DocTreeNode) -> dict[str, float]: - values = dict(node.confidence) - for child in node.children.values(): - for path, score in _collect_confidences(child).items(): - values[path] = max(values.get(path, 0.0), score) - return values - - -def _pop_leaf_path(node: DocTreeNode, path: str) -> bool: - if path in node.leaf_content: - node.leaf_content.pop(path) - return True - for child in node.children.values(): - if _pop_leaf_path(child, path): - return True - return False - - -def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int: - text = "\n".join(str(chunk.get("content") or "") for chunk in chunks) - return estimate_tokens(text) - - -async def _fetch_importance_norm_scores( - db: AsyncSession, - *, - user_id: str, - namespace: str, - chunk_ids: list[str], -) -> dict[str, float]: - if not chunk_ids: - return {} - stmt = ( - select( - RetrievalHitStat.chunk_id, - RetrievalHitStat.hit_count, - RetrievalHitStat.last_hit_at, - RetrievalHitStat.created_at, - ) - .where(RetrievalHitStat.user_id == user_id) - .where(RetrievalHitStat.namespace == namespace) - .where(RetrievalHitStat.hit_kind == "chunk") - .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) - ) - result = await db.execute(stmt) - scores: dict[str, float] = {} - for chunk_id, hit_count, last_hit_at, created_at in result.all(): - if chunk_id and last_hit_at and created_at: - scores[str(chunk_id)] = compute_importance_score( - hit_count, - last_hit_at, - created_at, - ) - return scores - - -async def trim_evidence_to_budget( - db: AsyncSession, - *, - doc_trees: dict[str, DocTreeNode], - doc_id_to_name: dict[str, str], - context_remaining: int, - user_id: str, - namespace: str, - ledger: BudgetLedger | None, - safety_margin: float = 0.9, -) -> str: - full_text = await render_evidence(db, doc_trees, doc_id_to_name) - target = int(max(context_remaining, 0) * safety_margin) - if estimate_tokens(full_text) <= target: - return full_text - - candidates: list[tuple[str, str, tuple[float, float, float], int]] = [] - for doc_id, tree in doc_trees.items(): - confidence = _collect_confidences(tree) - for path, chunks in _iter_leaf_content(tree): - chunk_ids = [ - str(chunk.get("chunk_id")) - for chunk in chunks - if chunk.get("chunk_id") - ] - importance = 0.0 - importance_scores = await _fetch_importance_norm_scores( - db, - user_id=user_id, - namespace=namespace, - chunk_ids=chunk_ids, - ) - if importance_scores: - importance = max(importance_scores.values()) - discovery_score = ( - float(chunks[0].get("discovery_score", 0.0) or 0.0) - if chunks - else 0.0 - ) - score = ( - float(confidence.get(path, 0.0) or 0.0), - discovery_score, - importance, - ) - candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks))) - - current_estimate = estimate_tokens(full_text) - removed: list[dict[str, Any]] = [] - for doc_id, path, score, token_estimate in sorted( - candidates, - key=lambda item: (item[2], -item[3]), - ): - if current_estimate <= target: - break - if _pop_leaf_path(doc_trees[doc_id], path): - confidence_score, discovery_score, importance_score = score - removed.append( - { - "document_id": doc_id, - "document_name": doc_id_to_name.get(doc_id, doc_id), - "path": path, - "confidence_score": round(confidence_score, 4), - "discovery_score": round(discovery_score, 4), - "importance_score": round(importance_score, 4), - "token_estimate": token_estimate, - } - ) - current_estimate = max(current_estimate - token_estimate, 0) - - if ledger is not None: - ledger.trimmed_paths.extend(removed) - logger.info( - f" agentic.trim_evidence: removed={len(removed)} " - f"est_tokens={current_estimate} target={target}" - ) - return await render_evidence(db, doc_trees, doc_id_to_name) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py deleted file mode 100644 index 3b8bd9c73..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py +++ /dev/null @@ -1,448 +0,0 @@ -"""Render agentic document trees into evidence text.""" -from __future__ import annotations - -from typing import Any, cast - -from shared.services.retrieval.agentic.core.types import DocTreeNode - -AssetLookupValue = str - - -def render_unified_doc_tree( - node: DocTreeNode, - doc_name: str, - depth: int = 0, - asset_lookup: dict[str, AssetLookupValue] | None = None, -) -> str: - """Render a DocTreeNode as one coherent hierarchy.""" - parts: list[str] = [] - indent = " " * depth - - if depth == 0: - parts.append(f"[Document] {doc_name}\n") - - child_prefixes = set(node.children.keys()) - - def min_sort(path: str) -> float: - chunks = node.leaf_content.get(path, []) - return min((chunk.get("sort_order") or float("inf") for chunk in chunks), default=float("inf")) - - render_queue: list[tuple[float, str, dict | str]] = [] - outline_paths: set[str] = set() - outline_position = 0.0 - - for item in node.outline_items: - path = item.get("path", "") - if any(path.startswith(child_prefix + " / ") for child_prefix in child_prefixes): - continue - outline_paths.add(path) - - if path in node.leaf_content or path in node.children: - sort_key = min_sort(path) if path in node.leaf_content else outline_position - else: - sort_key = outline_position - outline_position = max(outline_position, sort_key) + 0.001 - - render_queue.append((sort_key, "outline", item)) - - for path in node.leaf_content: - if path not in outline_paths: - render_queue.append((min_sort(path), "orphan_leaf", path)) - - for path in node.children: - if path not in outline_paths and path not in node.leaf_content: - child_sort = _infer_child_sort_order(node.children[path]) - render_queue.append((child_sort, "orphan_child", path)) - - render_queue.sort(key=lambda item: item[0]) - - for _sort_key, render_type, data in render_queue: - if render_type == "outline": - item = cast(dict, data) - path = item.get("path", "") - title = item.get("title", "") - is_leaf = item.get("is_leaf", False) - level = item.get("level", 1) - leaf_tag = " [Leaf]" if is_leaf else "" - - level_tag = f"[L{level}] " if level else "" - # Indent based on the section's own level relative to the tree depth, - # so L2 items are indented even when they sit in the root node. - item_indent = indent + " " * max(level - 1, 0) - if level <= 1: - parts.append(f"{item_indent}▸ {level_tag}{title}{leaf_tag}") - else: - parts.append(f"{item_indent}└ {level_tag}{title}{leaf_tag}") - - sub_indent = item_indent + " " - if path in node.children: - child = node.children[path] - if path in node.leaf_content: - # Parent has children: suppress page summary to avoid - # redundant parent/child overlap (§ SAME-AS pattern). - render_leaf_chunks( - parts, node.leaf_content[path], sub_indent, - asset_lookup=asset_lookup, suppress_page_summary=True, - ) - child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) - elif path in node.leaf_content: - render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - - elif render_type == "orphan_leaf": - path = cast(str, data) - title = path.rsplit(" / ", 1)[-1] if " / " in path else path - level = _infer_level_from_path(path) - item_indent = indent + " " * max(level - 1, 0) - sub_indent = item_indent + " " - level_tag = f"[L{level}] " if level else "" - if path in node.children: - # Non-leaf node with own content: render heading, then - # self chunks, then child subtree (merged rendering). - parts.append(f"{item_indent}▸ {level_tag}{title}") - render_leaf_chunks( - parts, node.leaf_content[path], sub_indent, - asset_lookup=asset_lookup, suppress_page_summary=True, - ) - child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) - else: - parts.append(f"{item_indent}▸ {level_tag}{title} [Leaf]") - render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - - elif render_type == "orphan_child": - path = cast(str, data) - title = path.rsplit(" / ", 1)[-1] if " / " in path else path - level = _infer_level_from_path(path) - item_indent = indent + " " * max(level - 1, 0) - level_tag = f"[L{level}] " if level else "" - child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) - # Only render the orphan heading if the child has content. - # Prevents empty orphan nodes from polluting evidence_text. - if child_text.strip(): - parts.append(f"{item_indent}▸ {level_tag}{title}") - parts.append(child_text) - - return "\n".join(parts) - - -def render_leaf_chunks( - parts: list[str], - chunks: list[dict[str, Any]], - indent: str, - asset_lookup: dict[str, AssetLookupValue] | None = None, - suppress_page_summary: bool = False, -) -> None: - chunk_by_id = { - chunk.get("chunk_id", ""): chunk - for chunk in chunks - if chunk.get("chunk_id") - } - rendered_ids: set[str] = set() - - for chunk in chunks: - chunk_id = chunk.get("chunk_id", "") - if chunk_id and chunk_id in rendered_ids: - continue - - chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "text").strip().lower() - if chunk_type in ("image", "table"): - continue - - if chunk_id: - rendered_ids.add(chunk_id) - - if chunk_type == "page": - render_page_chunk_lines( - parts, chunk, indent, - asset_lookup=asset_lookup, - suppress_summary=suppress_page_summary, - ) - continue - - content = str(chunk.get("content", "")).strip() - for connection in (chunk.get("chunk_metadata") or {}).get("connect_to") or []: - target = chunk_by_id.get(connection.get("target", "")) - if not target: - continue - target_id = target.get("chunk_id", "") - target_type = (target.get("chunk_type") or target.get("type") or "").strip().lower() - ref_str = connection.get("ref", "") - if not ref_str or ref_str not in content: - continue - - if target_id: - rendered_ids.add(target_id) - - if target_type == "table": - file_path = target.get("file_path") or "" - asset_url = _lookup_asset_url(asset_lookup, target_id) - display_ref = asset_url or file_path - table_lines = render_table_chunk_lines( - target, - display_ref=display_ref, - chunk_by_id=chunk_by_id, - asset_lookup=asset_lookup, - rendered_ids=rendered_ids, - ) - content = content.replace(ref_str, "\n" + "\n".join(table_lines) + "\n") - elif target_type == "image": - file_path = target.get("file_path") or "" - image_description = str(target.get("content", "")).strip() - if ref_str in image_description: - image_description = image_description.replace(ref_str, "").strip() - asset_url = _lookup_asset_url(asset_lookup, target_id) - display_ref = asset_url or file_path - if display_ref: - content = content.replace(ref_str, f"\n[Image: {display_ref}]\n{image_description}\n") - elif image_description: - content = content.replace(ref_str, f"\n[Image description]\n{image_description}\n") - - for line in content.split("\n"): - if line.strip(): - parts.append(f"{indent}┈ {line}") - - for chunk in chunks: - chunk_id = chunk.get("chunk_id", "") - if chunk_id and chunk_id in rendered_ids: - continue - if chunk_id: - rendered_ids.add(chunk_id) - - chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "").strip().lower() - if chunk_type == "image": - file_path = chunk.get("file_path") or "" - image_description = str(chunk.get("content", "")).strip() - asset_url = _lookup_asset_url(asset_lookup, chunk_id) - display_ref = asset_url or file_path - if display_ref: - parts.append(f"{indent}┈ [Image: {display_ref}]") - if image_description: - for line in image_description.split("\n"): - if line.strip(): - parts.append(f"{indent}┈ {line}") - elif chunk_type == "table": - file_path = chunk.get("file_path") or "" - asset_url = _lookup_asset_url(asset_lookup, chunk_id) - display_ref = asset_url or file_path - for line in render_table_chunk_lines( - chunk, - display_ref=display_ref, - chunk_by_id=chunk_by_id, - asset_lookup=asset_lookup, - rendered_ids=rendered_ids, - ): - if line.strip(): - parts.append(f"{indent}┈ {line}") - - -def render_page_chunk_lines( - parts: list[str], - chunk: dict[str, Any], - indent: str, - asset_lookup: dict[str, AssetLookupValue] | None = None, - suppress_summary: bool = False, -) -> None: - metadata = chunk.get("chunk_metadata") or chunk.get("metadata") or {} - summary = str(metadata.get("summary") or chunk.get("summary") or "").strip() - page_nums = _coerce_page_nums( - metadata.get("page_nums") or chunk.get("page_nums") - ) - if page_nums: - parts.append(f"{indent}┈ {_format_page_range(page_nums).capitalize()}") - if summary and not suppress_summary: - for line in summary.split("\n"): - if line.strip(): - parts.append(f"{indent}┈ {line}") - elif not page_nums: - parts.append(f"{indent}┈ [Page]") - - url = _asset_url_for_chunk(chunk, asset_lookup=asset_lookup) - if url: - parts.append(f"{indent}┈ [Page PDF ({_format_page_range(page_nums)}): {url}]") - - -def render_table_chunk_lines( - chunk: dict[str, Any], - *, - display_ref: str, - chunk_by_id: dict[str, dict[str, Any]] | None = None, - asset_lookup: dict[str, AssetLookupValue] | None = None, - rendered_ids: set[str] | None = None, -) -> list[str]: - header = f"[Table: {display_ref}]" if display_ref else "[Table]" - lines = [header] - table_path = chunk.get("source_chunk_path") or chunk.get("section_path") - if table_path: - lines.append(f"Table path: {table_path}") - file_path = chunk.get("file_path") - if file_path: - lines.append(f"Table asset: {file_path}") - - metadata = chunk.get("chunk_metadata") or chunk.get("metadata") or {} - summary = metadata.get("summary") if isinstance(metadata, dict) else "" - if summary: - lines.append("Table summary:") - lines.extend(str(summary).split("\n")) - - keywords = metadata.get("keywords") if isinstance(metadata, dict) else [] - if isinstance(keywords, list) and keywords: - keyword_text = ";".join(str(keyword) for keyword in keywords if str(keyword).strip()) - if keyword_text: - lines.append("Main columns:") - lines.append(keyword_text) - elif isinstance(keywords, str) and keywords.strip(): - lines.append("Main columns:") - lines.append(keywords.strip()) - - caption = metadata.get("caption") if isinstance(metadata, dict) else "" - if caption: - lines.append("Caption:") - lines.append(str(caption).strip()) - - lines.extend( - _render_table_embedded_image_lines( - chunk, - chunk_by_id=chunk_by_id or {}, - asset_lookup=asset_lookup, - rendered_ids=rendered_ids, - ) - ) - return lines - - -def _render_table_embedded_image_lines( - table_chunk: dict[str, Any], - *, - chunk_by_id: dict[str, dict[str, Any]], - asset_lookup: dict[str, AssetLookupValue] | None, - rendered_ids: set[str] | None, -) -> list[str]: - metadata = table_chunk.get("chunk_metadata") or table_chunk.get("metadata") or {} - if not isinstance(metadata, dict): - return [] - - lines: list[str] = [] - for connection in metadata.get("connect_to") or []: - if not isinstance(connection, dict): - continue - if str(connection.get("relation") or "").strip() != "embeds": - continue - target_id = str(connection.get("target") or "").strip() - if not target_id: - continue - target = chunk_by_id.get(target_id) - if not target: - continue - target_type = ( - target.get("chunk_type") or target.get("type") or "" - ).strip().lower() - if target_type != "image": - continue - if rendered_ids is not None: - rendered_ids.add(target_id) - - file_path = target.get("file_path") or "" - asset_url = _lookup_asset_url(asset_lookup, target_id) - display_ref = asset_url or file_path - image_description = str(target.get("content") or "").strip() - ref_str = str(connection.get("ref") or "").strip() - if ref_str and ref_str in image_description: - image_description = image_description.replace(ref_str, "").strip() - - if display_ref: - lines.append(f"[Image: {display_ref}]") - elif image_description: - lines.append("[Image description]") - if image_description: - lines.extend( - line for line in image_description.split("\n") if line.strip() - ) - return lines - - -def _asset_url_for_chunk( - chunk: dict[str, Any], - *, - asset_lookup: dict[str, AssetLookupValue] | None, -) -> str: - chunk_id = str(chunk.get("chunk_id") or "").strip() - value = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" - return str(value or "").strip() - - -def _lookup_asset_url( - asset_lookup: dict[str, AssetLookupValue] | None, - chunk_id: str, -) -> str: - if not chunk_id: - return "" - return str((asset_lookup or {}).get(chunk_id, "") or "").strip() - - -def _format_page_range(page_nums: list[int]) -> str: - pages = sorted(set(page_nums)) - if not pages: - return "pages unknown" - if len(pages) == 1: - return f"page {pages[0]}" - ranges: list[str] = [] - start = pages[0] - prev = pages[0] - for page in pages[1:]: - if page == prev + 1: - prev = page - continue - ranges.append(str(start) if start == prev else f"{start}-{prev}") - start = prev = page - ranges.append(str(start) if start == prev else f"{start}-{prev}") - return f"pages {', '.join(ranges)}" - - -def _coerce_page_nums(value: object) -> list[int]: - if isinstance(value, list): - raw_values = value - elif value is None: - raw_values = [] - else: - raw_values = str(value).split(",") - - pages: list[int] = [] - for item in raw_values: - try: - pages.append(int(str(item).strip())) - except (TypeError, ValueError): - continue - return pages - - -def _infer_level_from_path(path: str) -> int: - """Infer the section level for an orphan path from its segment count. - - Section paths use ``" / "`` as separator with one segment per hierarchy - level (e.g. ``"安全类 / SJSYJ-SC103 / ... / 表3 ..."`` → level 5). - This aligns with ``doc_nav.json`` level numbering. - """ - parts = [p for p in path.split(" / ") if p.strip()] - return max(len(parts), 1) - - -def _infer_child_sort_order(child: DocTreeNode) -> float: - """Infer sort position from the child's earliest chunk sort_order. - - When an orphan child node has no outline entry, we fall back to the - minimum ``sort_order`` across all its hydrated chunks so that orphans - render in document order instead of being appended at the end. - """ - min_order = float("inf") - for chunks in child.leaf_content.values(): - for chunk in chunks: - order = chunk.get("sort_order") - if order is not None and order < min_order: - min_order = float(order) - for grandchild in child.children.values(): - grandchild_order = _infer_child_sort_order(grandchild) - min_order = min(min_order, grandchild_order) - return min_order diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py deleted file mode 100644 index 8858056f4..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py +++ /dev/null @@ -1,713 +0,0 @@ -"""Legal action projection for agentic document navigation.""" -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Literal - -from shared.services.retrieval.agentic.core.budget import budget_status_from_snapshot -from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger -from shared.services.retrieval.agentic.navigation.state import ( - RejectionRecord, -) -from shared.services.retrieval.search.lexical_text import normalize_section_path -from shared.utils.text_utils import truncate_content_preview - - -ActionKind = Literal[ - "EXPAND", - "COLLECT", - "BACK", - "SEARCH_IMAGES", - "SEARCH_TABLES", - "FINISH", -] - - -@dataclass(frozen=True) -class LegalAction: - id: str - action: ActionKind - path: str | None = None - target_scope: str | None = None - asset_type: str | None = None - note: str | None = None - source: str = "tree" - score: float = 0.0 - critical_expand: bool = False - - -@dataclass -class LegalActionSet: - by_id: dict[str, LegalAction] = field(default_factory=dict) - expand: list[LegalAction] = field(default_factory=list) - collect: list[LegalAction] = field(default_factory=list) - back: list[LegalAction] = field(default_factory=list) - search: list[LegalAction] = field(default_factory=list) - finish: LegalAction | None = None - - def add(self, action: LegalAction) -> None: - self.by_id[action.id] = action - if action.action == "EXPAND": - self.expand.append(action) - elif action.action == "COLLECT": - self.collect.append(action) - elif action.action == "BACK": - self.back.append(action) - elif action.action in ("SEARCH_IMAGES", "SEARCH_TABLES"): - self.search.append(action) - elif action.action == "FINISH": - self.finish = action - - def get(self, action_id: str | None) -> LegalAction | None: - if not action_id: - return None - return self.by_id.get(action_id) - - -def build_legal_actions( - *, - items: list[dict[str, Any]], - current_scope: str | None, - collected_paths: list[dict[str, Any]], - expanded_scopes: set[str], - discovery_hints: list[dict[str, Any]] | None = None, - rejected: dict[str, RejectionRecord] | None = None, - total_images: int, - total_tables: int, - disabled_asset_types: set[str] | None = None, - budget_snapshot: dict[str, Any] | None = None, -) -> LegalActionSet: - action_set = LegalActionSet() - covered_paths = _covered_paths(collected_paths) - outline_paths = _outline_paths(collected_paths) - budget_mode = budget_status_from_snapshot(budget_snapshot) - rejection_ledger = rejected or {} - tool_adjudicated_paths = { - path for path, record in rejection_ledger.items() - if record.reason == "tool_adjudicated" - } - navigational_abandoned_paths = { - path for path, record in rejection_ledger.items() - if record.reason == "navigational_abandon" - } - discovery_scores = _discovery_scores_by_path(discovery_hints or []) - scored_items = _score_items(items, discovery_scores) - ranked_items = _rank_items(scored_items) - expand_allowlist = _expand_allowlist( - ranked_items, - budget_mode=budget_mode, - limit=3, - ) - - expand_index = 1 - collect_index = 1 - for item in scored_items: - path = str(item.get("path") or "").strip() - if not path or path == "Root": - continue - normalized_path = PathLedger.normalize(path) - if PathLedger.is_covered(normalized_path, covered_paths): - continue - # tool_adjudicated rejections are content-level negative and are not - # revived this round (TODO: future strong-signal revival). - if PathLedger.is_covered(normalized_path, tool_adjudicated_paths): - continue - - action_set.add(LegalAction( - id=f"C{collect_index}", - action="COLLECT", - path=path, - target_scope=path, - note=( - "upgrade outline to full evidence" - if normalized_path in outline_paths - else _item_note(item) - ), - score=float(item.get("relevance_score") or 0.0), - )) - collect_index += 1 - - critical_expand = False - if budget_mode == "EXHAUSTED": - continue - if budget_mode == "CRITICAL": - if action_set.collect: - continue - critical_expand = True - if item.get("is_leaf"): - continue - if normalized_path == current_scope: - continue - if current_scope and PathLedger.is_ancestor(normalized_path, current_scope): - continue - if normalized_path in expanded_scopes: - continue - if ( - budget_mode == "TIGHT" - and normalized_path not in expand_allowlist - ): - continue - # EXPAND suppression: navigational_abandon (weak) suppresses unless a - # discovery / lexical signal revives the path. - if normalized_path in navigational_abandoned_paths and not _has_discovery_signal( - normalized_path, discovery_scores - ): - continue - - action_set.add(LegalAction( - id=f"E{expand_index}", - action="EXPAND", - path=path, - target_scope=path, - note=_item_note(item), - score=float(item.get("relevance_score") or 0.0), - critical_expand=critical_expand, - )) - expand_index += 1 - - discovery_index = 1 - seen_discovery_paths: set[str] = set() - for hint in sorted( - discovery_hints or [], - key=lambda item: float(item.get("discovery_score") or 0.0), - reverse=True, - ): - path = normalize_section_path(str(hint.get("section_path") or "")) - if not path or path in seen_discovery_paths: - continue - seen_discovery_paths.add(path) - if PathLedger.is_covered(path, covered_paths): - continue - # tool_adjudicated rejections are not revived by discovery this round. - # TODO: allow tool-specific LLM adjudicators to revive rejected - # collects when validity cannot be determined structurally. - if PathLedger.is_covered(path, tool_adjudicated_paths): - continue - if any(action.path == path for action in action_set.collect): - continue - action_set.add(LegalAction( - id=f"D{discovery_index}", - action="COLLECT", - path=path, - target_scope=path, - note=_discovery_note(hint), - source="discovery", - score=float(hint.get("discovery_score") or 0.0), - )) - discovery_index += 1 - - search_allowed = budget_mode not in ("CRITICAL", "EXHAUSTED") - disabled_assets = {item.lower() for item in disabled_asset_types or set()} - if ( - search_allowed - and "image" not in disabled_assets - and total_images > 0 - and _asset_search_worthwhile(ranked_items, "image") - ): - action_set.add(LegalAction( - id="S1", - action="SEARCH_IMAGES", - asset_type="image", - note=f"{total_images} images available in current scope", - )) - if ( - search_allowed - and "table" not in disabled_assets - and total_tables > 0 - and _asset_search_worthwhile(ranked_items, "table") - ): - action_set.add(LegalAction( - id="S2", - action="SEARCH_TABLES", - asset_type="table", - note=f"{total_tables} tables available in current scope", - )) - - if current_scope and budget_mode != "EXHAUSTED": - back_index = 1 - for target in PathLedger.back_targets(current_scope): - label = target if target is not None else "root" - action_set.add(LegalAction( - id=f"B{back_index}", - action="BACK", - path=target, - target_scope=target, - note=f"return to {label}", - )) - back_index += 1 - - action_set.add(LegalAction( - id="F1", - action="FINISH", - note="finish this document", - )) - return action_set - - -def format_agent_state_block( - *, - current_scope: str | None, - query_intent: str, - expanded_scopes: set[str], - rejected: dict[str, RejectionRecord], - collected_paths: list[dict[str, Any]], - prior_tool_result: dict[str, Any] | None, - search_context: str, - budget_snapshot: dict[str, Any] | None, -) -> str: - lines = [ - "=== Agent State ===", - f"Current scope: {current_scope or 'root'}", - f"Advisory query intent: {query_intent or 'UNKNOWN'}", - _format_budget_state(budget_snapshot), - ] - budget_mode = budget_status_from_snapshot(budget_snapshot) - if budget_mode == "CRITICAL": - lines.append( - "Budget policy: exploration actions are closed; collect the best visible " - "evidence before FINISH." - ) - elif budget_mode == "EXHAUSTED": - lines.append( - "Budget policy: planning budget is exhausted or in overdraft. Do not " - "explore or search again. Use the current observation and tool results " - "to decide FINISH, or collect only indispensable visible evidence." - ) - if expanded_scopes: - lines.append("Expanded scopes:") - for path in sorted(expanded_scopes): - lines.append(f' - "{path}"') - else: - lines.append("Expanded scopes: none") - - navigational_abandoned = sorted( - path for path, record in rejected.items() - if record.reason == "navigational_abandon" - ) - tool_adjudicated = sorted( - path for path, record in rejected.items() - if record.reason == "tool_adjudicated" - ) - if navigational_abandoned: - lines.append( - "Scopes avoided (soft; revived by discovery):" - ) - for path in navigational_abandoned: - lines.append(f' - "{path}"') - if tool_adjudicated: - lines.append( - "Collects rejected by tool reconciliation (content-level; not revived):" - ) - for path in tool_adjudicated: - lines.append(f' - "{path}"') - - full_paths, outline_paths = _dedupe_collection_modes(collected_paths) - if full_paths or outline_paths: - if full_paths: - lines.append(f"Full evidence collected: {len(full_paths)} item(s)") - for path in full_paths: - lines.append(f' - "{path}"') - else: - lines.append("Full evidence collected: none") - if outline_paths: - lines.append( - f"Outline-only evidence: {len(outline_paths)} item(s) " - "(structure only; not hydrated as full chunks)" - ) - for path in outline_paths: - lines.append(f' - "{path}"') - else: - lines.append("Collected evidence: none") - - if prior_tool_result: - lines.append(f"Last tool result: {_compact_dict(prior_tool_result)}") - if search_context: - lines.append("Tool observation:") - lines.append(search_context.strip()) - lines.append("=== End Agent State ===") - return "\n".join(lines) - - -def format_actionable_observation( - *, - items: list[dict[str, Any]], - action_set: LegalActionSet, - max_chars: int = 20000, -) -> tuple[str, bool]: - """Render visible document state and legal action affordances once.""" - if not items: - return "(no visible sections)", False - - full_text = _render_actionable_items( - items=items, - action_set=action_set, - include_summary=True, - ) - if len(full_text) <= max_chars: - return full_text, False - - slim_text = _render_actionable_items( - items=items, - action_set=action_set, - include_summary=False, - ) - return slim_text[:max_chars], True - - -def _render_actionable_items( - *, - items: list[dict[str, Any]], - action_set: LegalActionSet, - include_summary: bool, -) -> str: - collect_by_path = { - action.path: action - for action in action_set.collect - if action.path - } - expand_by_path = { - action.path: action - for action in action_set.expand - if action.path - } - lines = [ - "=== Actionable Observation ===", - "Each visible section appears once. Choose action IDs attached to the relevant line.", - ] - for item in items: - lines.extend(_render_actionable_item( - item=item, - collect_action=collect_by_path.get(str(item.get("path") or "")), - expand_action=expand_by_path.get(str(item.get("path") or "")), - include_summary=include_summary, - )) - - discovery_lines = _format_discovery_actions(action_set) - if discovery_lines: - lines.append("") - lines.append("Discovery hints:") - lines.extend(discovery_lines) - - global_actions = _format_global_actions(action_set) - if global_actions: - lines.append("") - lines.append("Global actions:") - lines.extend(global_actions) - lines.append("=== End Actionable Observation ===") - return "\n".join(lines) - - -def _render_actionable_item( - *, - item: dict[str, Any], - collect_action: LegalAction | None, - expand_action: LegalAction | None, - include_summary: bool, -) -> list[str]: - level = int(item.get("level", 1) or 1) - show_summary = bool(item.get("show_summary", True)) - has_actions = collect_action is not None or expand_action is not None - show_details = show_summary or has_actions - path = str(item.get("path") or "") - summary = str(item.get("summary") or "") - is_leaf = bool(item.get("is_leaf", False)) - indent = " " * max(level - 1, 0) - prefix = "▸" if level == 1 else "└" - level_tag = f"depth={level}" - counts = _format_counts(item) if show_details else "" - tokens = _format_token_estimate(item) if show_details else "" - leaf = " [Leaf]" if is_leaf else "" - actions = _format_node_actions( - collect_action=collect_action, - expand_action=expand_action, - ) - - lines = [ - f'{indent}{prefix} {level_tag} path="{path}"{counts}{tokens}{leaf} actions: {actions}' - ] - if include_summary and show_details and summary: - display_summary = _enrich_section_covers_summary(summary) - clipped = truncate_content_preview(display_summary, head=120, tail=0) - lines.append(f"{indent} summary: {clipped}") - return lines - - -def _format_node_actions( - *, - collect_action: LegalAction | None, - expand_action: LegalAction | None, -) -> str: - actions: list[str] = [] - if collect_action: - collect_name = ( - "collect_full" - if collect_action.note == "upgrade outline to full evidence" - else "collect" - ) - actions.append(f"{collect_name}={collect_action.id}") - if expand_action: - actions.append(f"expand={expand_action.id}") - return ", ".join(actions) if actions else "none" - - -def _format_discovery_actions(action_set: LegalActionSet) -> list[str]: - lines: list[str] = [] - for action in action_set.collect: - if action.source != "discovery" or not action.path: - continue - note = f" | {action.note}" if action.note else "" - lines.append(f' {action.id} -> "{action.path}"{note}') - return lines - - -def _format_global_actions(action_set: LegalActionSet) -> list[str]: - lines: list[str] = [] - for action in action_set.search: - if action.action == "SEARCH_IMAGES": - lines.append(f" search_images={action.id} ({action.note})") - elif action.action == "SEARCH_TABLES": - lines.append(f" search_tables={action.id} ({action.note})") - for action in action_set.back: - target = action.target_scope or "root" - lines.append(f" back={action.id} -> {target}") - if action_set.finish: - lines.append(f" finish={action_set.finish.id}") - return lines - - -def _format_counts(item: dict[str, Any]) -> str: - parts: list[str] = [] - chunk_count = int(item.get("chunk_count") or 0) - image_count = int(item.get("image_count") or 0) - table_count = int(item.get("table_count") or 0) - if chunk_count: - parts.append(f"text={chunk_count}") - if image_count: - parts.append(f"image={image_count}") - if table_count: - parts.append(f"table={table_count}") - return f' [{" ".join(parts)}]' if parts else "" - - -def _format_token_estimate(item: dict[str, Any]) -> str: - total_chars = int(item.get("total_chars") or 0) - if total_chars <= 0: - return "" - tokens = total_chars / 2 - if tokens >= 1000: - return f" ~{tokens / 1000:.1f}k tokens" - return f" ~{int(tokens)} tokens" - - -def _enrich_section_covers_summary(summary: str) -> str: - prefix = "This section covers: " - if not summary.startswith(prefix): - return summary - body = summary[len(prefix):] - sub_sections = [s.strip() for s in body.split(", ") if s.strip()] - return f"This section covers {len(sub_sections)} sub-sections: {body}" - - -def _covered_paths(collected_paths: list[dict[str, Any]]) -> set[str]: - return { - PathLedger.normalize(str(item.get("path") or "")) - for item in collected_paths - if item.get("path") and item.get("hydrate_mode") != "outline" - } - - -def _outline_paths(collected_paths: list[dict[str, Any]]) -> set[str]: - return { - str(item.get("path") or "") - for item in collected_paths - if item.get("path") and item.get("hydrate_mode") == "outline" - } - - -def _dedupe_collection_modes( - collected_paths: list[dict[str, Any]], -) -> tuple[list[str], list[str]]: - full: set[str] = set() - outline: set[str] = set() - for item in collected_paths: - path = str(item.get("path") or "") - if not path: - continue - if item.get("hydrate_mode") == "outline": - outline.add(path) - else: - full.add(path) - outline -= full - return sorted(full), sorted(outline) - - -def _discovery_note(hint: dict[str, Any]) -> str | None: - summary = str(hint.get("summary") or "").strip() - score = float(hint.get("discovery_score") or 0.0) - score_note = f"score={score:.2f}" if score > 0 else "" - if summary: - clipped = truncate_content_preview(summary, head=120, tail=0) - return f"{clipped} {score_note}".strip() - chunk_type = str(hint.get("chunk_type") or "").strip() - if chunk_type: - return f"bottom-discovery hit type={chunk_type} {score_note}".strip() - return f"bottom-discovery hit {score_note}".strip() - - -def _format_budget_state(snapshot: dict[str, Any] | None) -> str: - if not isinstance(snapshot, dict): - return "Budget mode: UNKNOWN" - planning = snapshot.get("planning") - if not isinstance(planning, dict): - return "Budget mode: UNKNOWN" - status = str(planning.get("status") or "UNKNOWN") - used_pct = planning.get("used_pct") - remaining = planning.get("remaining") - capacity = planning.get("capacity") - overdraft = int(planning.get("overdraft") or 0) - overdraft_note = f", overdraft={overdraft}" if overdraft > 0 else "" - if used_pct is None: - return f"Budget mode: {status}" - if remaining is not None and capacity: - return ( - f"Budget mode: {status} ({used_pct}% used, " - f"{remaining}/{capacity} tokens remaining{overdraft_note})" - ) - return f"Budget mode: {status} ({used_pct}% used{overdraft_note})" - - -def _item_note(item: dict[str, Any]) -> str | None: - parts: list[str] = [] - chunk_count = int(item.get("chunk_count") or 0) - image_count = int(item.get("image_count") or 0) - table_count = int(item.get("table_count") or 0) - if chunk_count: - parts.append(f"text={chunk_count}") - if image_count: - parts.append(f"image={image_count}") - if table_count: - parts.append(f"table={table_count}") - score = float(item.get("relevance_score") or 0.0) - if score > 0: - parts.append(f"relevance={score:.2f}") - if item.get("is_leaf"): - parts.append("leaf") - return " ".join(parts) if parts else None - - -def _discovery_scores_by_path( - discovery_hints: list[dict[str, Any]], -) -> dict[str, float]: - scores: dict[str, float] = {} - for hint in discovery_hints: - path = normalize_section_path(str(hint.get("section_path") or "")) - if not path: - continue - score = float(hint.get("discovery_score") or 0.0) - scores[path] = max(scores.get(path, 0.0), score) - return scores - - -def _score_items( - items: list[dict[str, Any]], - discovery_scores: dict[str, float], -) -> list[dict[str, Any]]: - scored: list[dict[str, Any]] = [] - for index, item in enumerate(items): - copied = dict(item) - copied["_original_index"] = index - copied["relevance_score"] = _score_item(copied, discovery_scores) - scored.append(copied) - return scored - - -def _rank_items(scored_items: list[dict[str, Any]]) -> list[dict[str, Any]]: - return sorted( - scored_items, - key=lambda item: ( - float(item.get("relevance_score") or 0.0), - int(item.get("chunk_count") or 0), - int(item.get("table_count") or 0) + int(item.get("image_count") or 0), - -int(item.get("_original_index") or 0), - ), - reverse=True, - ) - - -def _score_item( - item: dict[str, Any], - discovery_scores: dict[str, float], -) -> float: - path = normalize_section_path(str(item.get("path") or "")) - if not path: - return 0.0 - score = discovery_scores.get(path, 0.0) - for hint_path, hint_score in discovery_scores.items(): - if PathLedger.is_ancestor(path, hint_path): - score = max(score, float(hint_score) * 0.9) - elif PathLedger.is_ancestor(hint_path, path): - score = max(score, float(hint_score) * 0.65) - return min(score, 1.0) - - -def _expand_allowlist( - ranked_items: list[dict[str, Any]], - *, - budget_mode: str, - limit: int, -) -> set[str]: - if budget_mode != "TIGHT": - return { - normalize_section_path(str(item.get("path") or "")) - for item in ranked_items - if item.get("path") - } - candidates = [ - normalize_section_path(str(item.get("path") or "")) - for item in ranked_items - if item.get("path") - and not item.get("is_leaf") - ] - return set(candidates[:limit]) - - -def _has_discovery_signal( - path: str, - discovery_scores: dict[str, float], -) -> bool: - """A path has a discovery signal if it, an ancestor, or a descendant appears in discovery.""" - return any( - candidate == path - or PathLedger.is_ancestor(path, candidate) - or PathLedger.is_ancestor(candidate, path) - for candidate in discovery_scores - ) - - -def _asset_search_worthwhile( - ranked_items: list[dict[str, Any]], - asset_kind: Literal["image", "table"], -) -> bool: - count_key = "image_count" if asset_kind == "image" else "table_count" - return any(int(item.get(count_key) or 0) > 0 for item in ranked_items[:5]) - - -def _compact_dict(value: dict[str, Any]) -> str: - bits: list[str] = [] - for key in ("tool", "status", "matched", "candidate_count", "status_detail"): - if key in value: - bits.append(f"{key}={value[key]}") - budget = value.get("budget") - if isinstance(budget, dict): - delta = budget.get("delta") - after = budget.get("after") - if isinstance(delta, dict): - bits.append( - "budget_delta=" - f"used:{delta.get('used', 0)}, " - f"used_pct:{delta.get('used_pct', 0)}, " - f"overdraft:{delta.get('overdraft', 0)}" - ) - if isinstance(after, dict) and int(after.get("overdraft") or 0) > 0: - bits.append(f"budget_overdraft={after.get('overdraft')}") - return ", ".join(bits) if bits else str(value) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py deleted file mode 100644 index 038f70001..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py +++ /dev/null @@ -1,818 +0,0 @@ -from __future__ import annotations - -import time -from typing import Any - -from loguru import logger -from sqlalchemy import func as sa_func -from sqlalchemy import or_, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import DocumentChunk, DocumentSection -from shared.models.database.job_result import JobResult -from shared.services.retrieval.agentic.core.budget import BudgetExceeded -from shared.services.retrieval.hydration.assets import build_retrieval_asset_url_map -from shared.services.retrieval.llm_adapter import LLMFn -from shared.utils.token_estimate import estimate_tokens - - -def build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]: - owner_map: dict[str, str] = {} - for chunk in text_chunks: - if (chunk.get("chunk_type") or "text") != "text": - continue - section_path = chunk.get("section_path") or "" - if not section_path: - continue - metadata = chunk.get("chunk_metadata") or {} - if not isinstance(metadata, dict): - continue - for conn in metadata.get("connect_to") or []: - if not isinstance(conn, dict): - continue - target_id = str(conn.get("target") or "").strip() - if target_id and target_id not in owner_map: - owner_map[target_id] = section_path - return owner_map - - -async def _load_scope_sections( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - scope_paths: list[str], -) -> list[tuple[str, str]]: - section_stmt = ( - select(DocumentSection.section_id, DocumentSection.section_path) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - if scope_paths: - scope_filters = [] - for scope in scope_paths: - scope_filters.append(DocumentSection.section_path == scope) - scope_filters.append(DocumentSection.section_path.like(f"{scope} / %")) - section_stmt = section_stmt.where(or_(*scope_filters)) - rows = (await db.execute(section_stmt)).all() - return [(section_id, section_path or "") for section_id, section_path in rows] - - -async def count_assets_under_scope( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - scope_paths: list[str], -) -> tuple[int, int]: - section_rows = await _load_scope_sections( - db, - document_id=document_id, - job_result_id=job_result_id, - scope_paths=scope_paths, - ) - all_section_ids = [section_id for section_id, _section_path in section_rows] - - if not all_section_ids: - return 0, 0 - - count_stmt = ( - select( - DocumentChunk.chunk_type, - sa_func.count(DocumentChunk.id), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .where(DocumentChunk.chunk_type.in_(["image", "table"])) - .group_by(DocumentChunk.chunk_type) - ) - count_result = await db.execute(count_stmt) - - total_images = 0 - total_tables = 0 - for chunk_type, count in count_result.all(): - if chunk_type == "image": - total_images = count - elif chunk_type == "table": - total_tables = count - return total_images, total_tables - - -async def resolve_root_asset_owners( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - chunks: list[dict[str, Any]], -) -> dict[str, str]: - root_asset_ids = [ - str(chunk.get("chunk_id") or "") - for chunk in chunks - if not chunk.get("owner_section_path") - and (chunk.get("section_path") or "") == "Root" - and (chunk.get("chunk_type") or "").lower() in ("image", "table") - and chunk.get("chunk_id") - ] - if not root_asset_ids: - return {} - - root_asset_set = set(root_asset_ids) - text_stmt = ( - select( - DocumentChunk.chunk_metadata, - DocumentSection.section_path, - ) - .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_type == "text") - ) - result = await db.execute(text_stmt) - - owner_map: dict[str, str] = {} - for metadata, section_path in result.all(): - if not isinstance(metadata, dict) or not section_path: - continue - for conn in metadata.get("connect_to") or []: - if not isinstance(conn, dict): - continue - target_id = str(conn.get("target") or "").strip() - if target_id in root_asset_set and target_id not in owner_map: - owner_map[target_id] = section_path - - if owner_map: - logger.info( - f" resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} " - f"Root assets to their owner sections" - ) - return owner_map - - -async def asset_filter_step( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - scope_path: str | list[str] | None, - asset_type: str, -) -> list[dict[str, Any]]: - t0 = time.monotonic() - try: - scope_list = ( - scope_path - if isinstance(scope_path, list) - else [scope_path] - if scope_path - else [] - ) - - section_rows = await _load_scope_sections( - db, - document_id=document_id, - job_result_id=job_result_id, - scope_paths=scope_list, - ) - section_ids = {row[0] for row in section_rows} - - if not section_ids: - logger.info(f" asset_filter_step: no sections found under scope={scope_path}") - return [] - - section_path_by_id = { - section_id: section_path for section_id, section_path in section_rows - } - asset_rows = ( - await db.execute( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.section_id, - DocumentChunk.source_chunk_path, - DocumentChunk.chunk_metadata, - DocumentChunk.sort_order, - DocumentChunk.job_result_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(section_ids))) - .where(DocumentChunk.chunk_type == asset_type) - .order_by(DocumentChunk.sort_order) - ) - ).all() - - text_rows = ( - await db.execute( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.chunk_metadata, - DocumentChunk.source_chunk_path, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(section_ids))) - .where(DocumentChunk.chunk_type == "text") - ) - ).all() - text_row_dicts = [ - { - "chunk_type": chunk_type, - "chunk_metadata": metadata or {}, - "section_id": section_id, - "section_path": section_path_by_id.get(section_id, ""), - "source_chunk_path": source_chunk_path, - } - for section_id, chunk_type, metadata, source_chunk_path in text_rows - ] - owner_by_target_id = build_connected_owner_map(text_row_dicts) - - connected_target_ids: set[str] = set(owner_by_target_id.keys()) - if connected_target_ids: - connected_rows = ( - await db.execute( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.section_id, - DocumentChunk.source_chunk_path, - DocumentChunk.chunk_metadata, - DocumentChunk.sort_order, - DocumentChunk.job_result_id, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(connected_target_ids))) - .where(DocumentChunk.chunk_type == asset_type) - .order_by(DocumentChunk.sort_order) - ) - ).all() - else: - connected_rows = [] - - job_id = ( - await db.execute(select(JobResult.job_id).where(JobResult.id == job_result_id)) - ).scalar() or "" - seen_ids: set[str] = set() - chunks: list[dict[str, Any]] = [] - for row in list(asset_rows) + list(connected_rows): - chunk_id = row[0] - if chunk_id in seen_ids: - continue - seen_ids.add(chunk_id) - - owner_section_path = owner_by_target_id.get(chunk_id) - if not owner_section_path: - own_section_path = section_path_by_id.get(row[4]) - if own_section_path and own_section_path == "Root": - logger.warning( - " asset_filter_step: rejecting root-level owner fallback " - f"chunk_id={chunk_id} section_path={own_section_path}" - ) - own_section_path = None - owner_section_path = own_section_path - - if not owner_section_path: - logger.warning( - f" asset_filter_step unresolved owner: chunk_id={chunk_id} " - f"file_path={row[3]} scope={scope_path or 'root'}" - ) - continue - - chunks.append( - { - "document_id": document_id, - "chunk_id": chunk_id, - "chunk_type": row[1], - "content": row[2], - "file_path": row[3], - "section_id": row[4], - "section_path": owner_section_path, - "owner_section_path": owner_section_path, - "source_chunk_path": row[5], - "chunk_metadata": row[6] or {}, - "sort_order": row[7], - "job_result_id": job_result_id, - "job_id": job_id, - } - ) - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" asset_filter_step scope={scope_path or 'root'} " - f"type={asset_type}: {len(chunks)} chunks found, {latency}ms" - ) - return chunks - - except Exception as exc: - logger.error(f" asset_filter_step failed: {exc}") - return [] - - -async def search_assets_step( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - scope_path: str | list[str] | None, - asset_type: str, - query: str, - llm_fn: LLMFn, - vlm_fn: LLMFn | None = None, -) -> dict[str, Any]: - """LLM-filtered asset search. - - For **tables**: uses text LLM with summary descriptions (unchanged). - For **images**: generates presigned S3 URLs via - ``build_retrieval_asset_url_map`` and sends them to the VLM - (``vlm_fn``) for visual relevance judgment. - """ - t0 = time.monotonic() - - all_assets = await asset_filter_step( - db, - document_id=document_id, - job_result_id=job_result_id, - scope_path=scope_path, - asset_type=asset_type, - ) - if not all_assets: - logger.info(f" search_assets_step: no {asset_type} assets under scope={scope_path}") - return { - "status": "empty", - "matched_assets": [], - "verdicts": [], - "candidate_count": 0, - } - - # Build lookup by chunk_id - asset_by_id: dict[str, dict[str, Any]] = {} - for asset in all_assets: - chunk_id = str(asset.get("chunk_id") or "") - if chunk_id: - asset_by_id[chunk_id] = asset - - if not asset_by_id: - return { - "status": "empty", - "matched_assets": [], - "verdicts": [], - "candidate_count": 0, - } - - status_detail = "" - status = "empty" - - # ── Route by asset type ────────────────────────────────────────── - if asset_type == "image": - if vlm_fn is None: - logger.info(" search_assets_step: VLM unavailable for image search") - selected_ids = await _search_assets_via_text_llm( - query=query, - asset_type=asset_type, - assets=list(asset_by_id.values()), - llm_fn=llm_fn, - ) - status = "fallback_matched" if selected_ids else "fallback_empty" - status_detail = "vlm_unavailable_text_fallback" - else: - selected_ids, vlm_error = await _search_images_via_vlm( - query=query, - assets=list(asset_by_id.values()), - vlm_fn=vlm_fn, - ) - if vlm_error: - logger.info( - " search_assets_step: VLM image search fell back to text " - f"filter, reason={vlm_error}" - ) - selected_ids = await _search_assets_via_text_llm( - query=query, - asset_type=asset_type, - assets=list(asset_by_id.values()), - llm_fn=llm_fn, - ) - status = "fallback_matched" if selected_ids else "fallback_empty" - status_detail = "vlm_failed_text_fallback" - else: - status = "matched" if selected_ids else "empty" - else: - selected_ids = await _search_assets_via_text_llm( - query=query, - asset_type=asset_type, - assets=list(asset_by_id.values()), - llm_fn=llm_fn, - ) - status = "matched" if selected_ids else "empty" - - selected_id_set = {str(cid) for cid in selected_ids} - matched_assets = [asset_by_id[cid] for cid in selected_ids if cid in asset_by_id] - verdicts = [ - _asset_verdict( - asset, - relevant=str(asset.get("chunk_id") or "") in selected_id_set, - reason=( - _selected_reason(status) - if str(asset.get("chunk_id") or "") in selected_id_set - else _not_selected_reason(status) - ), - ) - for asset in asset_by_id.values() - ] - - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f" search_assets_step query=\"{query}\" type={asset_type}: " - f"{len(matched_assets)}/{len(all_assets)} assets matched, {latency}ms" - ) - return { - "status": status, - "status_detail": status_detail, - "matched_assets": matched_assets, - "verdicts": verdicts, - "candidate_count": len(asset_by_id), - "latency_ms": latency, - } - - -def _asset_verdict( - asset: dict[str, Any], - *, - relevant: bool, - reason: str, -) -> dict[str, Any]: - metadata = asset.get("chunk_metadata") or {} - summary = metadata.get("summary", "") - return { - "chunk_id": asset.get("chunk_id", ""), - "file_path": asset.get("file_path", ""), - "section_path": asset.get("owner_section_path") or asset.get("section_path", ""), - "summary": summary, - "relevant": relevant, - "reason": reason, - } - - -def _selected_reason(status: str) -> str: - if status.startswith("fallback_"): - return "selected_by_text_fallback" - return "selected_by_asset_inspector" - - -def _not_selected_reason(status: str) -> str: - if status.startswith("fallback_"): - return "not_selected_by_text_fallback" - return "not_selected_by_asset_inspector" - - -async def _search_assets_via_text_llm( - *, - query: str, - asset_type: str, - assets: list[dict[str, Any]], - llm_fn: LLMFn, -) -> list[str]: - """Text-based LLM filtering for table assets.""" - candidates_for_llm, valid_ids, id_to_chunk_id = _project_assets_for_text_filter( - query=query, - asset_type=asset_type, - assets=assets, - ) - - prompt = _format_asset_filter_prompt(query, asset_type, candidates_for_llm) - try: - response = await llm_fn(prompt) - selected_ids = _parse_asset_filter_response(response, valid_ids) - return [ - id_to_chunk_id[row_id] - for row_id in selected_ids - if row_id in id_to_chunk_id - ] - except BudgetExceeded: - raise - except Exception as exc: - logger.warning(f" _search_assets_via_text_llm failed: {exc}") - return [] - - -def _project_assets_for_text_filter( - *, - query: str, - asset_type: str, - assets: list[dict[str, Any]], -) -> tuple[list[dict[str, str]], set[str], dict[str, str]]: - """Project assets into a prompt-sized text view. - - Stable row identifiers are shown to the model. Owner paths stay internal: - reconciliation and hydration use the original asset rows, not prompt text. - Descriptive text is reduced only when the complete prompt would exceed the - navigation planning budget envelope. - """ - 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}" - metadata = asset.get("chunk_metadata") or {} - summary = str(metadata.get("summary") or "").strip() - file_path = str(asset.get("file_path") or "") - content = str(asset.get("content") or "").strip() - description = summary 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(query, asset_type, projected) - prompt_budget = _asset_filter_prompt_budget() - if estimate_tokens(prompt) <= prompt_budget: - return projected, valid_ids, id_to_chunk_id - - structural_prompt = _format_asset_filter_prompt( - query, - asset_type, - [ - { - "id": item["id"], - "file": item["file"], - "desc": "", - } - for item in projected - ], - ) - structural_tokens = estimate_tokens(structural_prompt) - desc_budget = max(prompt_budget - structural_tokens, len(projected)) - per_item_desc_tokens = max(desc_budget // len(projected), 1) - compacted = [ - { - "id": item["id"], - "file": item["file"], - "desc": _fit_text_to_token_budget(item["desc"], per_item_desc_tokens), - } - for item in projected - ] - return compacted, valid_ids, id_to_chunk_id - - -def _asset_filter_prompt_budget() -> int: - from shared.services.retrieval.agentic.core.runtime import build_config_from_env - - config = build_config_from_env() - planning_capacity = int( - max(config.token_budget_total - config.bootstrap_budget, 0) - * config.planning_ratio - ) - return max(planning_capacity, 1) - - -def _fit_text_to_token_budget(text: str, token_budget: int) -> str: - text = text.strip() - if not text or estimate_tokens(text) <= token_budget: - return text - words = text.split() - if len(words) > 1: - kept: list[str] = [] - for word in words: - candidate = " ".join([*kept, word]) - if estimate_tokens(candidate) > token_budget: - break - kept.append(word) - return " ".join(kept).strip() - - lo = 0 - hi = len(text) - best = "" - while lo <= hi: - mid = (lo + hi) // 2 - candidate = text[:mid].strip() - if estimate_tokens(candidate) <= token_budget: - best = candidate - lo = mid + 1 - else: - hi = mid - 1 - return best - - -async def _search_images_via_vlm( - *, - query: str, - assets: list[dict[str, Any]], - vlm_fn: LLMFn, -) -> tuple[list[str], str | None]: - """VLM-based image search with presigned S3 URLs. - - Generates presigned URLs for each image asset, builds a multimodal - prompt with image_url blocks, and asks the VLM to select relevant ones. - """ - url_map = await build_retrieval_asset_url_map( - assets, log_context="search_images_vlm", - ) - - # Only include images that have valid URLs. - candidates: list[tuple[str, str, str]] = [] # (row_id, file_path, url) - 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 = url_map.get(chunk_id) - if not url or isinstance(url, list): - continue - row_id = f"I{index}" - file_path = asset.get("file_path") or "" - candidates.append((row_id, file_path, url)) - valid_ids.add(row_id) - id_to_chunk_id[row_id] = chunk_id - - if not candidates: - logger.info(" _search_images_via_vlm: no presigned URLs available, skipping") - return [], "no_presigned_urls" - - messages = _format_vlm_image_filter_messages(query, candidates) - try: - response = await vlm_fn(messages) - selected_ids = _parse_asset_filter_response(response, valid_ids) - return [ - id_to_chunk_id[row_id] - for row_id in selected_ids - if row_id in id_to_chunk_id - ], None - except BudgetExceeded: - raise - except Exception as exc: - logger.warning(f" _search_images_via_vlm failed: {exc}") - return [], str(exc) - - -def _format_asset_filter_prompt( - query: str, - asset_type: str, - candidates: list[dict[str, str]], -) -> str: - """Build the text LLM prompt for table asset filtering.""" - 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 _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 _markdown_cell(value: str) -> str: - return ( - str(value or "") - .replace("\n", " ") - .replace("\r", " ") - .replace("|", "\\|") - .strip() - ) - - -def _format_vlm_image_filter_messages( - query: str, - candidates: list[tuple[str, str, str]], -) -> list[dict[str, Any]]: - """Build multimodal VLM messages with inline image URLs. - - Each candidate is (row_id, file_path, presigned_url). - The VLM sees the actual images and decides relevance. - """ - 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"Selection policy:\n" - f"- Match both the requested visual type and requested subject.\n" - f"- Do not select images only because they are charts or from " - f"the same broad domain.\n" - f"- Do not broaden specific market, instrument, company, metric, " - f"or entity terms. Neighboring topics are not matches unless " - f"the image explicitly connects them to the requested subject.\n" - f"- Treat words like \"all\" as all relevant images, not all " - f"visible candidates.\n" - f"- If uncertain, do not select the image.\n\n" - ), - }, - ] - - for row_id, file_path, url in candidates: - content_parts.append({ - "type": "text", - "text": f'Image {row_id} file="{file_path}":', - }) - content_parts.append({ - "type": "image_url", - "image_url": {"url": url}, - }) - - content_parts.append({ - "type": "text", - "text": ( - f"\n\nReturn ONLY a JSON array of matching image row IDs, e.g.: " - f'["{candidates[0][0]}"]\n' - f"If none are relevant, return an empty array: []\n" - f"Do not include any explanation." - ), - }) - - return [{"role": "user", "content": content_parts}] - - -def _parse_asset_filter_response( - text: str, - valid_ids: set[str], -) -> list[str]: - """Parse LLM response for asset filter and keep valid row IDs.""" - import json - import re - - text = text.strip() - - # Try direct JSON parse - 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 - - # Try extracting from code fence - 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 - - # Try finding any JSON array - 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 [] diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py deleted file mode 100644 index 22d9c2736..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py +++ /dev/null @@ -1,1337 +0,0 @@ -"""Per-document navigation for agentic retrieval. - -Collector Agent architecture -~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The navigation loop uses an observe-act Collector Agent model where each step -produces one main action and optional collection side effects: - -1. **collect**: paths to add to the evidence collection -2. **action**: EXPAND/BACK/SEARCH_IMAGES/SEARCH_TABLES/FINISH - -The ``collected_paths`` list accumulates across all steps. After -navigation completes (or is interrupted), a single batch hydration -pass loads content for all collected paths. - -Asset collection (images/tables) still runs during navigation so LLM -tool requests are honoured, but assets are reconciled after hydration. -""" -from __future__ import annotations - -from typing import Any, cast - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.agentic import tools -from shared.services.retrieval.agentic.core.budget import BudgetExceeded -from shared.services.retrieval.agentic.evidence.builder import reconcile_deferred_assets -from shared.services.retrieval.agentic.core.runtime import AgentLlmBudget -from shared.services.retrieval.agentic.core.trace import TraceRecorder -from shared.services.retrieval.agentic.core.types import ( - AgentRunConfig, - AgentState, - CandidateDoc, - DecisionTraceStep, - DocTreeNode, - NavigateStepResult, - ToolResult, -) -from shared.services.retrieval.agentic.navigation.selection_hydration import ( - hydrate_path_selections_into_node, -) -from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger -from shared.services.retrieval.agentic.navigation.state import NavigationState -from shared.services.retrieval.agentic.prompts import ( - QUERY_INTENT_PROMPT, - parse_query_intent_response, -) -from shared.services.retrieval.llm_adapter import LLMFn - - -class DocumentNavigationRunner: - def __init__( - self, - *, - db: AsyncSession, - state: AgentState, - trace: TraceRecorder, - trace_enabled: bool, - user_id: str, - namespace: str, - query: str, - config: AgentRunConfig, - discovery_by_doc: dict[str, list[dict[str, Any]]], - llm_fn: LLMFn | None, - llm_budget: AgentLlmBudget, - disabled_asset_types: set[str] | None = None, - ) -> None: - self._db = db - self._state = state - self._trace = trace - self._trace_enabled = trace_enabled - self._user_id = user_id - self._namespace = namespace - self._query = query - self._config = config - self._discovery_by_doc = discovery_by_doc - self._llm_fn = llm_fn - self._llm_budget = llm_budget - self._disabled_asset_types = disabled_asset_types or set() - self._decision_steps: list[dict[str, Any]] = [] - - @property - def decision_steps(self) -> list[dict[str, Any]]: - return list(self._decision_steps) - - async def navigate_selected_documents(self) -> None: - logger.info( - f" agentic: Phase 2 — navigating {len(self._state.selected_docs)} documents" - ) - query_intent = await self._classify_query_intent() - for doc in self._state.selected_docs: - if self._state.elapsed_ms >= self._config.latency_budget_ms: - logger.info(" agentic: latency budget hit during Phase 2, stopping") - break - await self._navigate_document(doc, query_intent=query_intent) - - async def _navigate_document( - self, - doc: CandidateDoc, - *, - query_intent: str, - ) -> None: - job_result_id = self._state.doc_job_map.get(doc.document_id, "") - if not job_result_id: - logger.info(f" agentic: skipping doc {doc.document_id} — no job_result_id") - self._state.ever_explored_doc_ids.add(doc.document_id) - return - - doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "") - root = DocTreeNode(scope_path=None) - doc_pending_assets: list[dict[str, Any]] = [] - from shared.services.retrieval.agentic.navigation.section_tree import ( - load_document_section_rows, - ) - section_rows = await load_document_section_rows( - self._db, - document_id=doc.document_id, - job_result_id=job_result_id, - ) - - # Phase 2A: Collector Agent navigation (summary-only, no content hydration) - doc_pending_assets, collected_paths = await self._navigate_collector( - doc=doc, - root=root, - doc_name=doc_name, - job_result_id=job_result_id, - section_rows=section_rows, - query_intent=query_intent, - ) - - # Phase 2B: Batch hydrate all collected paths - if collected_paths: - await self._hydrate_collected( - doc=doc, - root=root, - job_result_id=job_result_id, - collected_paths=collected_paths, - section_rows=section_rows, - ) - - # Phase 2C: Reconcile assets into hydrated tree - if doc_pending_assets: - self._reconcile_pending_assets( - doc=doc, - root=root, - doc_name=doc_name, - doc_pending_assets=doc_pending_assets, - ) - - if doc.document_id in self._state.doc_trees: - self._state.doc_trees[doc.document_id].merge(root) - else: - self._state.doc_trees[doc.document_id] = root - self._state.ever_explored_doc_ids.add(doc.document_id) - if self._state.ledger is not None: - self._state.ledger.mark_explored(docs=1) - - async def _navigate_collector( - self, - *, - doc: CandidateDoc, - root: DocTreeNode, - doc_name: str, - job_result_id: str, - section_rows: list, - query_intent: str, - ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Collector Agent navigation loop. - - Returns (doc_pending_assets, collected_paths). - """ - doc_discovery_hints = self._discovery_by_doc.get(doc.document_id, []) - doc_pending_assets: list[dict[str, Any]] = [] - nav_state = NavigationState( - document_id=doc.document_id, - document_name=doc_name, - job_result_id=job_result_id, - ) - # Context from SEARCH tools — injected into next navigate prompt - search_context: str = "" - prior_tool_result: dict[str, Any] | None = None - - exit_reason = "unknown" - budget_failure: dict[str, Any] | None = None - - while nav_state.step_count < self._config.max_nav_steps: - if self._state.elapsed_ms >= self._config.latency_budget_ms: - exit_reason = "latency" - break - has_tool_context = bool(prior_tool_result or search_context) - if ( - self._state.ledger - and self._state.ledger.status("planning") == "EXHAUSTED" - and not has_tool_context - ): - logger.info(" agentic: planning budget exhausted, ending navigation for current doc") - exit_reason = "budget" - break - - nav_state.step_count += 1 - before_scope = nav_state.current_scope - expanded_before = set(nav_state.expanded_scopes) - rejected_before = dict(nav_state.rejected) - collected_before_count = len(nav_state.collected_paths) - - doc_llm_fn = self._llm_budget.for_document( - cast(LLMFn, self._llm_fn), - doc_id=doc.document_id, - step=nav_state.step_count, - allow_overdraft=has_tool_context, - overdraft_reason=( - "report_tool_result_to_main_agent" - if has_tool_context else "" - ), - ) - try: - nav_result = await tools.navigate_step( - self._db, - document_id=doc.document_id, - job_result_id=job_result_id, - query=self._query, - llm_fn=doc_llm_fn, - user_id=self._user_id, - namespace=self._namespace, - doc_name=doc_name, - scope_path=nav_state.current_scope, - budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, - nav_trace=nav_state.nav_trace if nav_state.nav_trace else None, - collected_paths=nav_state.collected_paths, - expanded_scopes=nav_state.expanded_scopes, - rejected=nav_state.rejected, - disabled_asset_types=self._disabled_asset_types | nav_state.blocked_asset_types_for_scope( - nav_state.current_scope - ), - discovery_hints=doc_discovery_hints, - section_rows=section_rows, - query_intent=query_intent, - search_context=search_context, - prior_tool_result=prior_tool_result, - ) - except BudgetExceeded as exc: - budget_failure = getattr(exc, "details", {}) or {} - logger.info( - " agentic: planning budget exhausted during navigation " - f"details={budget_failure}" - ) - if self._trace_enabled: - self._trace.record_budget_stop("planning_exhausted") - exit_reason = "budget" - break - self._state.step_count += 1 - - # Clear previous tool context (consumed by this step's prompt) - search_context = "" - prior_tool_result = None - - # ── Execute asset tools (SEARCH) ───────────────────────────── - asset_tool_result = await self._execute_asset_tools( - doc=doc, - job_result_id=job_result_id, - scope=nav_state.current_scope, - nav_result=nav_result, - pending_assets=doc_pending_assets, - parent_step_index=len(self._decision_steps), - ) - search_context = asset_tool_result.get("context", "") - prior_tool_result = asset_tool_result.get("summary") - if ( - prior_tool_result is not None - and self._should_block_asset_search( - prior_tool_result, nav_state.current_scope, - ) - ): - nav_state.block_asset_search( - nav_state.current_scope, - str(prior_tool_result.get("asset_type") or ""), - ) - - # Merge outline + confidence into root tree - _merge_step_node(root, nav_result.node) - - requested_collects = list(nav_result.collect) - collect_reconcile = self._reconcile_collects_after_tool( - nav_result=nav_result, - asset_tool_result=asset_tool_result, - ) - rejected_collects = collect_reconcile["rejected_collects"] - if rejected_collects: - nav_result.collect = collect_reconcile["accepted_collects"] - for path in rejected_collects: - nav_state.mark_rejected_collect( - path, - step=nav_state.step_count, - detail=collect_reconcile.get("reason", ""), - ) - logger.info( - " agentic: tool reconciliation rejected collects: " - f"{rejected_collects}" - ) - - # ── Process COLLECT ────────────────────────────────────────── - collected_in_step: list[str] = [] - for coll_item in nav_result.collect: - path = coll_item["path"] - nav_state.add_collected( - coll_item, - step=nav_state.step_count, - scope_context=nav_state.current_scope, - ) - collected_in_step.append(path) - # Outline collections keep children visible — the intent is - # "see structure, then drill deeper for full content". - # Coverage / action filtering now derives from collected_paths - # via the navigation state ledger (no physical exclusion). - - # ── Process navigation action ──────────────────────────────── - should_break = False - if ( - nav_result.action == "EXPAND" - and nav_result.result_status == "ok" - and nav_result.drill_into - ): - drill_path = nav_result.drill_into - # Create child node in tree for the drill target - target_parent = _find_target_node(root, drill_path) - target_parent.children.setdefault(drill_path, DocTreeNode(scope_path=drill_path)) - nav_state.mark_expanded(drill_path) - nav_state.current_scope = drill_path - - elif nav_result.action == "BACK" and nav_result.result_status == "ok": - if nav_state.current_scope is None: - logger.info(" agentic: BACK at root scope, staying at root") - nav_result.result_status = "invalid_back" - nav_result.result_note = "already_at_root" - else: - back_target = nav_result.back_to # None = root - if PathLedger.valid_back_target(nav_state.current_scope, back_target): - nav_state.mark_rejected_if_unproductive( - nav_state.current_scope, - step=nav_state.step_count, - detail="back_from_unproductive_scope", - ) - nav_state.current_scope = back_target - else: - logger.warning( - f" agentic: invalid back_to='{back_target}' " - f"from scope='{nav_state.current_scope}'" - ) - nav_result.result_status = "invalid_back" - nav_result.result_note = f"invalid_back_target: {back_target}" - - elif nav_result.action == "ERROR": - logger.warning( - f" agentic: navigation ERROR for doc={doc.document_id}: " - f"{nav_result.error_reason or nav_result.reason}" - ) - exit_reason = "error" - should_break = True - - elif nav_result.action == "FINISH" and nav_result.result_status == "ok": - exit_reason = "llm_finish" - should_break = True - - # ── Build trace entry ──────────────────────────────────────── - state_delta = nav_state.snapshot_delta( - before_scope=before_scope, - expanded_before=expanded_before, - rejected_before=rejected_before, - collected_before_count=collected_before_count, - ) - trace_entry: dict[str, Any] = { - "step": nav_state.step_count, - "scope": before_scope or "root", - "action": nav_result.action, - "drill_into": nav_result.drill_into, - "back_to": nav_result.back_to, - "collected": collected_in_step, - "tools_used": nav_result.tools, - "reason": nav_result.reason, - "result_status": nav_result.result_status, - "state_delta": state_delta, - } - if rejected_collects: - trace_entry["requested_collects"] = [ - item.get("path", "") - for item in requested_collects - if item.get("path") - ] - trace_entry["rejected_collects"] = rejected_collects - trace_entry["tool_reconciliation"] = collect_reconcile["reason"] - # Record tool usage & results so future steps can see search history - if prior_tool_result: - trace_entry["tool_results"] = prior_tool_result - nav_state.nav_trace.append(trace_entry) - - # ── Record decision step ───────────────────────────────────── - main_step_index = self._record_navigation_step( - doc=doc, - scope=before_scope, - step_num=nav_state.step_count, - nav_result=nav_result, - collected_in_step=collected_in_step, - asset_summary=asset_tool_result.get("summary"), - rejected_collects=rejected_collects, - state_delta=state_delta, - ) - asset_trace = asset_tool_result.get("asset_trace") - if asset_trace: - asset_trace["parent_step_index"] = main_step_index - self._append_decision_trace_step(DecisionTraceStep(**asset_trace)) - - if should_break: - break - else: - # while loop exhausted — max_nav_steps reached - exit_reason = "max_steps" - - # Hard guard: if navigation is forcefully interrupted and collected - # nothing, collect visible leaf children under the last explored - # scope. Voluntary FINISH/BACK with empty collection is respected. - # budget – planning pool EXHAUSTED (pre-check or exception) - # latency – elapsed time exceeded latency budget - # max_steps – navigation step count limit reached - # error – unexpected exception during navigate_step - forced_exits = ("budget", "latency", "max_steps", "error") - guard_triggered = False - if not nav_state.collected_paths and nav_state.step_count > 0 and exit_reason in forced_exits: - guard_triggered = True - guard_scope = nav_state.current_scope - logger.info( - f" agentic: forced exit ({exit_reason}) with 0 collected paths, " - f"auto-collecting leaves under scope={guard_scope or 'root'}" - ) - from shared.services.retrieval.agentic.navigation.section_tree import ( - load_child_sections, - ) - guard_items = await load_child_sections( - self._db, - doc.document_id, - job_result_id, - guard_scope, - section_rows=section_rows, - ) - for item in guard_items: - if not item.get("show_summary", True): - continue - if item.get("is_leaf"): - path = item["path"] - nav_state.collected_paths.append({ - "path": path, - "confidence": 0.4, - "hydrate_mode": "chunks", - "collected_at_step": nav_state.step_count, - "scope_context": guard_scope or "root", - "guard_reason": f"forced_exit_{exit_reason}", - }) - - self._append_decision_trace_step(DecisionTraceStep( - step_index=len(self._decision_steps), - agent="system", - phase="system_guard", - document_id=doc.document_id, - document=doc.source_file_name or "", - scope=guard_scope or "root", - observation={ - "exit_reason": exit_reason, - "collected_count": 0, - }, - decision={ - "action": "auto_collect_visible_leaves", - "args": {"scope": guard_scope or "root"}, - "reason": "hard navigation constraint stopped the loop before evidence collection", - }, - result={ - "status": "guard_auto_collect", - "collected": [ - { - "path": p["path"], - "confidence": p.get("confidence", 0.0), - } - for p in nav_state.collected_paths - ], - "note": f"forced_exit_{exit_reason}", - }, - budget=self._state.ledger.snapshot() if self._state.ledger else None, - elapsed_ms=self._state.elapsed_ms, - )) - - # ── Navigate summary — record exit reason and final state ───── - doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "") - self._append_decision_trace_step(DecisionTraceStep( - step_index=len(self._decision_steps), - agent="navigator", - phase="navigate_summary", - document_id=doc.document_id, - document=doc_name, - scope=nav_state.current_scope or "root", - observation={ - "total_steps": nav_state.step_count, - "collected_count": len(nav_state.collected_paths), - "guard_triggered": guard_triggered, - "query_intent": query_intent, - "budget_failure": budget_failure, - }, - decision={ - "action": "summarize_navigation", - "args": {}, - "reason": f"Navigation ended with exit_reason={exit_reason}", - }, - result={ - "status": exit_reason, - "final_scope": nav_state.current_scope or "root", - "collected": [ - { - "path": p.get("path", ""), - "confidence": p.get("confidence", 0.0), - } - for p in nav_state.collected_paths - ], - }, - budget=self._state.ledger.snapshot() if self._state.ledger else None, - elapsed_ms=self._state.elapsed_ms, - )) - - return doc_pending_assets, nav_state.collected_paths - - async def _classify_query_intent( - self, - ) -> str: - """Classify query intent as advisory state. Fail-open to UNKNOWN.""" - if self._llm_fn is None: - return "UNKNOWN" - prompt = QUERY_INTENT_PROMPT.format( - query=self._query, - ) - try: - response = await self._llm_budget.call( - cast(LLMFn, self._llm_fn), - prompt, - pool="planning", - priority="low", - ) - except Exception as exc: - logger.info( - f" agentic: query intent classifier failed-open: {exc}" - ) - return "UNKNOWN" - return parse_query_intent_response(response) - - async def _hydrate_collected( - self, - *, - doc: CandidateDoc, - root: DocTreeNode, - job_result_id: str, - collected_paths: list[dict[str, Any]], - section_rows: list, - ) -> None: - """Batch-hydrate all collected paths after navigation completes.""" - if not collected_paths: - return - - # Deduplicate: keep the most complete evidence mode per path. A later - # full collect is an upgrade over an earlier outline collect even when - # the outline confidence was higher. - deduped: dict[str, dict[str, Any]] = {} - for item in collected_paths: - path = item["path"] - if path not in deduped or _collect_rank(item) > _collect_rank(deduped[path]): - deduped[path] = item - unique_selections = list(deduped.values()) - - # Ensure child nodes exist for each collected path so reparent can - # correctly route descendant chunks into the right subtree. - for item in unique_selections: - path = item["path"] - _ensure_child_node(root, path) - - await hydrate_path_selections_into_node( - self._db, - node=root, - path_selections=unique_selections, - user_id=self._user_id, - namespace=self._namespace, - document_id=doc.document_id, - job_result_id=job_result_id, - ) - - # Single reparent pass — tree structure is final. - root.reparent_leaf_content() - - # Load section tree outline for each collected child node, then - # build a proper sub-tree so the renderer can nest L3 under L2 etc. - from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections - for item in unique_selections: - path = item["path"] - child_node = root.children.get(path) - if child_node is None or child_node.outline_items: - continue # Skip if no child or already has outline - try: - section_items = await load_child_sections( - self._db, doc.document_id, job_result_id, path, - limit_depth=False, - section_rows=section_rows, - ) - if section_items: - # Filter out the scope node itself AND ancestor/sibling - # items. load_child_sections returns ancestor context - # for navigation prompts, but for evidence rendering the - # child node only needs its own descendants. - child_node.outline_items = [ - si for si in section_items - if si.get("path") != path - and si.get("path", "").startswith(path + " / ") - ] - # Build sub-tree from outline hierarchy and re-reparent - # so chunks are correctly nested (e.g. L3 under L2). - _build_outline_subtree(child_node) - except Exception as exc: - logger.warning(f" hydrate_collected: failed to load outline for '{path}': {exc}") - - # Also organize any discovery chunks sitting in root.leaf_content. - # Same path-based tree building; pre-existing children are protected. - if root.leaf_content: - _build_outline_subtree(root) - - # Accurate budget accounting — count only actually-hydrated chunks. - if self._state.ledger is not None: - total_chunks = len(root.flatten_chunk_rows()) - self._state.ledger.mark_explored(chunks=total_chunks) - logger.info( - f" agentic: hydrate_collected doc={doc.document_id} " - f"collected={len(unique_selections)} hydrated_chunks={total_chunks}" - ) - - async def _execute_asset_tools( - self, - *, - doc: CandidateDoc, - job_result_id: str, - scope: str | None, - nav_result: NavigateStepResult, - pending_assets: list[dict[str, Any]], - parent_step_index: int, - ) -> dict[str, Any]: - """Execute SEARCH_* and return the observation for the next loop.""" - empty_result: dict[str, Any] = { - "context": "", - "summary": None, - "asset_trace": None, - } - if not nav_result.tools: - return empty_result - - if not nav_result.search_assets_params: - return empty_result - - params = nav_result.search_assets_params - search_query = params["query"] - asset_type = params["asset_type"] - scope_paths = params.get("scope_paths") - tool_name = "SEARCH_IMAGES" if asset_type == "image" else "SEARCH_TABLES" - budget_before = self._state.ledger.snapshot() if self._state.ledger else None - - search_llm_fn = self._llm_budget.for_document( - cast(LLMFn, self._llm_fn), - doc_id=doc.document_id, - step=self._state.step_count, - allow_overdraft=True, - overdraft_reason=f"{tool_name}_asset_inspector", - ) if self._llm_fn else None - if search_llm_fn is None: - summary = { - "tool": tool_name, - "asset_type": asset_type, - "query": search_query, - "matched": 0, - "status": "unavailable", - "scope_paths": scope_paths if scope_paths is not None else ([scope] if scope else []), - "matched_paths": [], - "sub_agent_assessment": "LLM unavailable for asset inspection", - } - return { - "context": self._format_asset_context( - tool_name, - asset_type, - search_query, - [], - status="unavailable", - ), - "summary": summary, - "asset_trace": self._build_asset_trace_payload( - doc=doc, - scope=scope, - parent_step_index=parent_step_index, - asset_type=asset_type, - query=search_query, - candidates=[], - result={"status": "unavailable", "verdicts": [], "matched_assets": []}, - ), - } - - vlm_fn: LLMFn | None = None - if asset_type == "image": - from shared.services.retrieval.llm_adapter import create_retrieval_vlm_fn - raw_vlm_fn = create_retrieval_vlm_fn() - if raw_vlm_fn is not None: - vlm_fn = self._llm_budget.for_document( - raw_vlm_fn, - doc_id=doc.document_id, - step=self._state.step_count, - allow_overdraft=True, - overdraft_reason=f"{tool_name}_vlm_asset_inspector", - ) - - try: - asset_result = await tools.search_assets_step( - self._db, - document_id=doc.document_id, - job_result_id=job_result_id, - scope_path=scope_paths if scope_paths is not None else scope, - asset_type=asset_type, - query=search_query, - llm_fn=search_llm_fn, - vlm_fn=vlm_fn, - ) - except BudgetExceeded as exc: - failure = getattr(exc, "details", {}) or {} - logger.info( - f" agentic: {tool_name} skipped — planning budget exhausted " - f"details={failure}" - ) - asset_result = { - "status": "budget_exceeded", - "status_detail": "budget_reserve_failed", - "budget_failure": failure, - "matched_assets": [], - "verdicts": [], - "candidate_count": 0, - } - - budget_after = self._state.ledger.snapshot() if self._state.ledger else None - budget_delta = _budget_delta(budget_before, budget_after) - matched_assets = asset_result.get("matched_assets") or [] - verdicts = asset_result.get("verdicts") or [] - candidate_count = int(asset_result.get("candidate_count") or 0) - if matched_assets: - pending_assets.extend(matched_assets) - - summary = { - "tool": tool_name, - "asset_type": asset_type, - "query": search_query, - "matched": len(matched_assets), - "candidate_count": candidate_count, - "status": asset_result.get("status", "empty"), - "status_detail": asset_result.get("status_detail", ""), - "budget": { - "before": _compact_budget_snapshot(budget_before), - "after": _compact_budget_snapshot(budget_after), - "delta": budget_delta, - }, - "scope_paths": scope_paths if scope_paths is not None else ([scope] if scope else []), - "matched_paths": [ - asset.get("file_path", "") - for asset in matched_assets - if asset.get("file_path") - ], - "matched_owner_paths": [ - asset.get("owner_section_path") or asset.get("section_path") or "" - for asset in matched_assets - if asset.get("owner_section_path") or asset.get("section_path") - ], - "sub_agent_assessment": ( - f"asset inspector matched {len(matched_assets)} " - f"of {candidate_count} {asset_type} candidates " - f"(status={asset_result.get('status', 'empty')})" - ), - } - logger.info( - f" agentic step {self._state.step_count}: {tool_name} " - f'doc="{doc.source_file_name}" scope={scope or "root"} ' - f'search_scope={scope_paths if scope_paths is not None else scope or "root"} ' - f'query="{search_query}" matched={len(matched_assets)}' - ) - return { - "context": self._format_asset_context( - tool_name, - asset_type, - search_query, - matched_assets, - status=str(asset_result.get("status", "empty")), - status_detail=str(asset_result.get("status_detail", "")), - ), - "summary": summary, - "matched_assets": matched_assets, - "candidate_count": candidate_count, - "asset_trace": self._build_asset_trace_payload( - doc=doc, - scope=scope, - parent_step_index=parent_step_index, - asset_type=asset_type, - query=search_query, - candidates=verdicts, - result=asset_result, - budget_before=budget_before, - budget_after=budget_after, - budget_delta=budget_delta, - ), - } - - @staticmethod - def _should_block_asset_search( - summary: dict[str, Any] | None, - current_scope: str | None, - ) -> bool: - if not summary: - return False - if int(summary.get("matched") or 0) > 0: - return False - if not _tool_searched_current_scope(summary, current_scope): - return False - return str(summary.get("status") or "").lower() in { - "empty", - "fallback_empty", - "unavailable", - "error", - "budget_exceeded", - } - - @staticmethod - def _reconcile_collects_after_tool( - *, - nav_result: NavigateStepResult, - asset_tool_result: dict[str, Any], - ) -> dict[str, Any]: - accepted_collects = list(nav_result.collect) - empty = { - "accepted_collects": accepted_collects, - "rejected_collects": [], - "reason": "", - } - if not accepted_collects or not nav_result.tools: - return empty - summary = asset_tool_result.get("summary") - if not isinstance(summary, dict): - return empty - tool_name = str(summary.get("tool") or "") - if tool_name not in {"SEARCH_IMAGES", "SEARCH_TABLES"}: - return empty - status = str(summary.get("status") or "").lower() - if status not in {"empty", "fallback_empty", "matched", "fallback_matched"}: - return empty - - matched_assets = asset_tool_result.get("matched_assets") or [] - matched_owner_paths = [ - str(asset.get("owner_section_path") or asset.get("section_path") or "") - for asset in matched_assets - if asset.get("owner_section_path") or asset.get("section_path") - ] - still_accepted: list[dict[str, Any]] = [] - rejected_paths: list[str] = [] - for item in accepted_collects: - path = str(item.get("path") or "") - if not path: - continue - has_matching_asset = any( - PathLedger.is_same_or_descendant(owner_path, path) - for owner_path in matched_owner_paths - ) - if has_matching_asset: - still_accepted.append(item) - else: - rejected_paths.append(path) - - if not rejected_paths: - return empty - reason = ( - f"{tool_name} returned no valid matching assets under rejected " - f"collect paths; status={status}, matched={len(matched_assets)}" - ) - return { - "accepted_collects": still_accepted, - "rejected_collects": rejected_paths, - "reason": reason, - } - - def _reconcile_pending_assets( - self, - *, - doc: CandidateDoc, - root: DocTreeNode, - doc_name: str, - doc_pending_assets: list[dict[str, Any]], - ) -> None: - if doc_name and not root.children and not any( - item.get("path") == doc_name for item in root.outline_items - ): - root.outline_items.insert(0, {"path": doc_name, "level": 0}) - reconcile_deferred_assets(root, doc_pending_assets) - if self._trace_enabled: - self._trace.record_step( - "deferred_asset_reconcile", - ToolResult( - status="reconciled", - payload={ - "document_id": doc.document_id, - "pending_count": len(doc_pending_assets), - "placed_count": sum( - 1 for asset in doc_pending_assets - if str(asset.get("chunk_id") or "") in { - str(row.get("chunk_id") or "") - for row in root.flatten_chunk_rows() - } - ), - }, - ), - decision_reason=f"deferred_reconcile_{doc.source_file_name}", - ) - - @staticmethod - def _format_asset_context( - tool_name: str, - asset_type: str, - search_query: str, - matched_assets: list[dict[str, Any]], - *, - status: str = "empty", - status_detail: str = "", - ) -> str: - 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): - file_path = asset.get("file_path", "") - lines.append(f" {i + 1}. {file_path}") - owner_paths = _unique_asset_owner_paths(matched_assets) - if owner_paths: - lines.append("Owner sections with matching assets:") - for owner_path in owner_paths: - lines.append(f' - "{owner_path}"') - lines.append( - "Use these asset results and owner sections to decide collect, " - "finish, back, or further navigation." - ) - lines.append(f"=== End {tool_name} Results ===") - return "\n".join(lines) - - def _build_asset_trace_payload( - self, - *, - doc: CandidateDoc, - scope: str | None, - parent_step_index: int, - asset_type: str, - query: str, - candidates: list[dict[str, Any]], - result: dict[str, Any], - budget_before: dict[str, Any] | None = None, - budget_after: dict[str, Any] | None = None, - budget_delta: dict[str, Any] | None = None, - ) -> dict[str, Any]: - matched_assets = result.get("matched_assets") or [] - matched = [ - { - "chunk_id": asset.get("chunk_id", ""), - "file_path": asset.get("file_path", ""), - "section_path": asset.get("owner_section_path") - or asset.get("section_path", ""), - } - for asset in matched_assets - ] - return { - "step_index": len(self._decision_steps), - "agent": "asset_inspector", - "parent_step_index": parent_step_index, - "phase": "asset_inspect", - "document_id": doc.document_id, - "document": doc.source_file_name or "", - "scope": scope or "root", - "observation": { - "asset_type": asset_type, - "query": query, - "candidates": candidates, - }, - "decision": { - "action": "inspect_assets", - "args": {"asset_type": asset_type, "query": query}, - "reason": "judged each candidate against the requested evidence", - }, - "result": { - "status": result.get("status", "empty"), - "status_detail": result.get("status_detail", ""), - "verdicts": result.get("verdicts") or [], - "matched": matched, - "budget_failure": result.get("budget_failure"), - "budget": { - "before": _compact_budget_snapshot(budget_before), - "after": _compact_budget_snapshot(budget_after), - "delta": budget_delta or {}, - }, - }, - "budget": self._state.ledger.snapshot() if self._state.ledger else None, - "elapsed_ms": self._state.elapsed_ms, - } - - def _append_decision_trace_step(self, step: DecisionTraceStep) -> int: - step.step_index = len(self._decision_steps) - self._decision_steps.append(step.to_dict()) - if self._trace_enabled: - self._trace.record_decision_trace_step(step) - return step.step_index - - def _record_navigation_step( - self, - *, - doc: CandidateDoc, - scope: str | None, - step_num: int, - nav_result: NavigateStepResult, - collected_in_step: list[str], - asset_summary: dict[str, Any] | None = None, - rejected_collects: list[str] | None = None, - state_delta: dict[str, Any] | None = None, - ) -> int: - action = nav_result.action - reason = nav_result.reason - drill_into = nav_result.drill_into - doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "") - collected = [ - { - "path": item.get("path", ""), - "confidence": item.get("confidence", 0.0), - "hydrate_mode": item.get("hydrate_mode", "chunks"), - } - for item in nav_result.collect - ] - decision_args: dict[str, Any] = {} - if drill_into: - decision_args["target"] = drill_into - if action == "BACK": - decision_args["target"] = nav_result.back_to - if nav_result.search_assets_params: - decision_args["query"] = nav_result.search_assets_params.get("query", "") - decision_args["asset_type"] = nav_result.search_assets_params.get("asset_type", "") - projected_scope = scope or "root" - if action == "EXPAND" and drill_into: - projected_scope = drill_into - elif action == "BACK" and nav_result.result_status == "ok": - projected_scope = nav_result.back_to or "root" - - result_payload: dict[str, Any] = { - "status": nav_result.result_status, - "collected": collected, - "new_scope": projected_scope, - "note": nav_result.result_note, - } - if state_delta is not None: - result_payload["state_delta"] = state_delta - if rejected_collects: - result_payload["rejected_collects"] = rejected_collects - if nav_result.error_reason: - result_payload["error"] = nav_result.error_reason - if asset_summary: - result_payload["matched_assets"] = asset_summary.get("matched", 0) - result_payload["tool_status"] = asset_summary.get("status") - result_payload["tool_budget"] = asset_summary.get("budget") - result_payload["sub_agent_assessment"] = asset_summary.get( - "sub_agent_assessment" - ) - - trace_step = DecisionTraceStep( - step_index=len(self._decision_steps), - agent="navigator", - phase="navigate", - document=doc_name, - document_id=doc.document_id, - scope=scope or "root", - observation=nav_result.observation, - decision={ - "action": action, - "args": decision_args, - "reason": reason, - }, - result=result_payload, - budget=self._state.ledger.snapshot() if self._state.ledger else None, - elapsed_ms=self._state.elapsed_ms, - ) - step_index = self._append_decision_trace_step(trace_step) - - status_tag = ( - f" status={nav_result.result_status}" - if nav_result.result_status != "ok" - else "" - ) - scope_log = scope or "root" - logger.info( - f" agentic step {self._state.step_count}: navigate_step " - f'doc="{doc.source_file_name}" scope={scope_log} ' - f"step={step_num} action={action} tools={nav_result.tools} " - f'reason="{reason[:80]}" ' - f"collected={len(collected_in_step)} " - f"drill_into={drill_into} " - f"outline={len(nav_result.node.outline_items)}" - f"{status_tag}" - ) - return step_index - - -def _tool_searched_current_scope( - summary: dict[str, Any], - current_scope: str | None, -) -> bool: - current = PathLedger.normalize(current_scope) or "root" - scope_paths = summary.get("scope_paths") - if not isinstance(scope_paths, list) or not scope_paths: - return current == "root" - searched = [ - PathLedger.normalize(str(path or "")) - for path in scope_paths - if PathLedger.normalize(str(path or "")) - ] - if current == "root": - return not searched - return searched == [current] - - -def _compact_budget_snapshot(snapshot: dict[str, Any] | None) -> dict[str, Any]: - if not isinstance(snapshot, dict): - return {} - planning = snapshot.get("planning") - if not isinstance(planning, dict): - return {} - compact = { - "status": planning.get("status"), - "used_pct": planning.get("used_pct"), - "remaining": planning.get("remaining"), - "capacity": planning.get("capacity"), - "overdraft": planning.get("overdraft", 0), - } - overdraft_events = snapshot.get("overdraft_events") - if isinstance(overdraft_events, list) and overdraft_events: - compact["overdraft_events"] = overdraft_events[-3:] - return compact - - -def _budget_delta( - before: dict[str, Any] | None, - after: dict[str, Any] | None, -) -> dict[str, Any]: - before_planning = ( - before.get("planning") if isinstance(before, dict) else None - ) - after_planning = ( - after.get("planning") if isinstance(after, dict) else None - ) - if not isinstance(before_planning, dict) or not isinstance(after_planning, dict): - return {} - return { - "used": int(after_planning.get("used") or 0) - - int(before_planning.get("used") or 0), - "used_pct": int(after_planning.get("used_pct") or 0) - - int(before_planning.get("used_pct") or 0), - "remaining": int(after_planning.get("remaining") or 0) - - int(before_planning.get("remaining") or 0), - "overdraft": int(after_planning.get("overdraft") or 0) - - int(before_planning.get("overdraft") or 0), - } - - -def _unique_asset_owner_paths(matched_assets: list[dict[str, Any]]) -> list[str]: - seen: set[str] = set() - owner_paths: list[str] = [] - for asset in matched_assets: - owner_path = str(asset.get("owner_section_path") or asset.get("section_path") or "") - if not owner_path or owner_path in seen: - continue - seen.add(owner_path) - owner_paths.append(owner_path) - return owner_paths - - -def _find_target_node(node: DocTreeNode, path: str) -> DocTreeNode: - """Walk the tree to find the deepest existing node that owns *path*. - - Only recurse when *path* is a true descendant of a child (prefix match). - An exact match means the item belongs to the section itself, which is - managed by the *parent* node — the renderer already handles the case - where a path appears in both ``children`` and ``leaf_content``. - """ - for child_path, child in node.children.items(): - if PathLedger.is_ancestor(child_path, path): - return _find_target_node(child, path) - return node - - -def _merge_step_node(root: DocTreeNode, step_node: DocTreeNode) -> None: - """Route outline items and confidence from *step_node* to correct tree positions.""" - for item in step_node.outline_items: - path = item.get("path", "") - target = _find_target_node(root, path) - existing = {i.get("path") for i in target.outline_items} - if path not in existing: - target.outline_items.append(item) - - for path, conf in step_node.confidence.items(): - target = _find_target_node(root, path) - target.confidence[path] = max(target.confidence.get(path, 0), conf) - - -def _collect_rank(item: dict[str, Any]) -> tuple[int, float]: - mode = str(item.get("hydrate_mode") or "chunks") - mode_rank = 0 if mode == "outline" else 1 - confidence = float(item.get("confidence") or 0.0) - return (mode_rank, confidence) - - -def _ensure_child_node(root: DocTreeNode, path: str) -> None: - """Create an intermediate child node for *path* if it doesn't exist. - - When a non-leaf section is COLLECTed (e.g. "五、施工安全保证措施"), - hydration loads all descendant chunks (e.g. "五、... / 3.监控量测措施 / 3.1..."). - These chunks are initially placed in root.leaf_content. - ``reparent_leaf_content`` then moves them into the correct child subtree — - but only if a child node exists for the collected path. - - This function creates that child node so reparent can work correctly. - """ - # Don't create a child for root-level or if already exists - if not path: - return - - # Walk to find the deepest existing ancestor node - target = root - for child_path, child in root.children.items(): - if PathLedger.is_ancestor(child_path, path): - target = child - break - if path == child_path: - return # Already exists - - # Create the child node on the target - if path not in target.children: - target.children[path] = DocTreeNode(scope_path=path) - - -def _build_outline_subtree(node: DocTreeNode) -> None: - """Recursively build child nodes from outline_items + leaf_content paths. - - Every section_path is ``" / "``-separated; parent-child is prefix match. - This function creates children one level at a time, reparents chunks - into them, then recurses so the next level is handled correctly. - """ - if not node.outline_items and not node.leaf_content: - return - - # All known paths: outline metadata + actual chunk paths. - all_paths = ( - {item["path"] for item in node.outline_items} - | set(node.leaf_content.keys()) - ) - - # Find paths that have at least one descendant. - parent_paths: set[str] = set() - for path in all_paths: - for other in all_paths: - if other != path and PathLedger.is_ancestor(path, other): - parent_paths.add(path) - break - - if not parent_paths: - return # All items are leaves — nothing to nest. - - # KEY FIX: only keep top-level parents. If "A / B" and "A / B / C" - # are both parents, only create "A / B" as a child NOW. "A / B / C" - # will be created when the function recurses into "A / B". - parent_paths = { - pp for pp in parent_paths - if not any(pp != other and PathLedger.is_ancestor(other, pp) for other in parent_paths) - } - - # Track children that already exist (e.g. from collected hydration). - # Their outline_items are already populated — don't push duplicates. - pre_existing = set(node.children.keys()) - - # Create child nodes for top-level parent paths only. - for path in parent_paths: - if path not in node.children: - node.children[path] = DocTreeNode(scope_path=path) - - # Reparent existing children that are descendants of newly-created - # parents. E.g. if node.children already has "A / B" and we just - # created "A", move "A / B" under "A". This prevents "A / B" from - # appearing as an orphan_child at the wrong tree depth. - for parent_path in parent_paths: - if parent_path in pre_existing: - continue # Don't reparent into pre-existing nodes - parent_node = node.children[parent_path] - to_move = [ - cp for cp in list(node.children.keys()) - if cp != parent_path - and PathLedger.is_ancestor(parent_path, cp) - ] - for cp in to_move: - parent_node.children[cp] = node.children.pop(cp) - - # Split outline_items: keep items at this level, move descendants - # into newly-created children only (skip pre-existing ones). - kept: list[dict] = [] - for item in node.outline_items: - item_path = item["path"] - best_parent: str | None = None - for pp in parent_paths: - if PathLedger.is_ancestor(pp, item_path): - if best_parent is None or len(pp) > len(best_parent): - best_parent = pp - if best_parent and best_parent not in pre_existing: - node.children[best_parent].outline_items.append(item) - else: - kept.append(item) - node.outline_items = kept - - # Reparent FIRST so children receive their leaf_content, - # THEN recurse so deeper levels can be built from that content. - node.reparent_leaf_content() - for child in node.children.values(): - _build_outline_subtree(child) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py deleted file mode 100644 index c0196be66..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Knowledge-map overview for agentic document selection.""" -from __future__ import annotations - -from sqlalchemy import func, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document, DocumentChunk, GraphNode - - -_MAX_OVERVIEW_FILES = 50 - - -async def build_knowledge_map_overview( - db: AsyncSession, - *, - user_id: str, - namespace: str, -) -> tuple[str, dict[str, str]]: - """Build a file-level knowledge map overview for LLM file selection.""" - doc_stmt = ( - select(Document) - .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.updated_at.desc()) - .limit(_MAX_OVERVIEW_FILES) - ) - doc_result = await db.execute(doc_stmt) - documents = list(doc_result.scalars()) - - if not documents: - return "(empty)", {} - - doc_ids = [document.document_id for document in documents] - doc_id_to_name = { - document.document_id: (document.source_file_name or document.document_id) - for document in documents - } - - chunk_stats_stmt = ( - select( - DocumentChunk.document_id, - func.count(DocumentChunk.id).label("chunk_count"), - func.count(func.nullif(DocumentChunk.chunk_type, "text")).label("media_count"), - ) - .join( - Document, - (Document.document_id == DocumentChunk.document_id) - & (Document.current_job_result_id == DocumentChunk.job_result_id), - ) - .where(DocumentChunk.document_id.in_(doc_ids)) - .group_by(DocumentChunk.document_id) - ) - chunk_stats_result = await db.execute(chunk_stats_stmt) - chunk_stats: dict[str, dict[str, int]] = {} - for document_id, chunk_count, media_count in chunk_stats_result.all(): - chunk_stats[document_id] = {"total": chunk_count, "media": media_count} - - graph_summary_stmt = ( - select(GraphNode.owner_document_id, GraphNode.properties) - .where(GraphNode.owner_document_id.in_(doc_ids)) - .where(GraphNode.node_kind == "document") - ) - graph_summary_result = await db.execute(graph_summary_stmt) - doc_top_summaries: dict[str, str] = {} - for document_id, properties in graph_summary_result.all(): - if not isinstance(properties, dict): - continue - top_summary = str(properties.get("top_summary") or "").strip() - if top_summary: - doc_top_summaries[document_id] = top_summary - - lines: list[str] = [] - for document in documents: - document_id = document.document_id - name = doc_id_to_name[document_id] - stats = chunk_stats.get(document_id, {"total": 0, "media": 0}) - top_summary = doc_top_summaries.get(document_id, "") - - line = f'- [{document_id}] {name} chunks={stats["total"]}' - if stats["media"] > 0: - line += f' media={stats["media"]}' - if top_summary: - line += f"\n top_summary:\n{indent_block(top_summary, 4)}" - lines.append(line) - - return "\n".join(lines), doc_id_to_name - - -def indent_block(text: str, spaces: int) -> str: - prefix = " " * spaces - return "\n".join(f"{prefix}{line}" for line in str(text or "").splitlines()) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/path_ledger.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/path_ledger.py deleted file mode 100644 index bd051eb5a..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/path_ledger.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Path relationship helpers for document navigation state.""" -from __future__ import annotations - -from collections.abc import Iterable - -from shared.services.retrieval.search.lexical_text import normalize_section_path - - -class PathLedger: - """Small, authoritative wrapper for section path relationships.""" - - @staticmethod - def normalize(path: str | None) -> str: - return normalize_section_path(str(path or "").strip()) - - @classmethod - def is_ancestor(cls, ancestor: str | None, descendant: str | None) -> bool: - ancestor_path = cls.normalize(ancestor) - descendant_path = cls.normalize(descendant) - if not ancestor_path or not descendant_path: - return False - return descendant_path.startswith(ancestor_path + " / ") - - @classmethod - def is_same_or_descendant(cls, path: str | None, scope: str | None) -> bool: - candidate = cls.normalize(path) - scope_path = cls.normalize(scope) - if not candidate or not scope_path: - return False - return candidate == scope_path or candidate.startswith(scope_path + " / ") - - @classmethod - def is_covered(cls, path: str | None, covered_paths: Iterable[str]) -> bool: - candidate = cls.normalize(path) - if not candidate: - return False - return any( - candidate == covered - or candidate.startswith(covered + " / ") - for covered in (cls.normalize(item) for item in covered_paths) - if covered - ) - - @classmethod - def back_targets(cls, current_scope: str | None) -> list[str | None]: - scope = cls.normalize(current_scope) - if not scope: - return [] - parts = [part for part in scope.split(" / ") if part] - targets: list[str | None] = [ - " / ".join(parts[:index]) - for index in range(len(parts) - 1, 0, -1) - ] - targets.append(None) - return targets - - @classmethod - def valid_back_target(cls, current_scope: str | None, target: str | None) -> bool: - scope = cls.normalize(current_scope) - if not scope: - return False - return target is None or cls.is_ancestor(target, scope) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py deleted file mode 100644 index 2e7ae2373..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Section count aggregation for agentic navigation.""" -from __future__ import annotations - -from sqlalchemy import case, func, literal_column, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import DocumentChunk - - -async def attach_section_counts( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - all_sections: dict[str, dict], - items_by_path: dict[str, dict], -) -> None: - """Attach direct chunk and connected asset counts to visible section items.""" - scope_item_sids = { - item["section_id"] - for item in items_by_path.values() - if item["show_summary"] - } - all_section_ids = [meta["section_id"] for meta in all_sections.values()] - if not all_section_ids or not scope_item_sids: - return - - section_id_counts = await _load_direct_chunk_counts( - db, - document_id=document_id, - job_result_id=job_result_id, - all_section_ids=all_section_ids, - ) - - sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()} - for section_id, (text_count, image_count, table_count, total_chars) in section_id_counts.items(): - chunk_path = sid_to_path.get(section_id, "") - if not chunk_path: - continue - - for item_path, item in items_by_path.items(): - if not item["show_summary"]: - continue - if _chunk_belongs_to_item(chunk_path, item_path): - item["chunk_count"] += text_count - item["image_count"] += image_count - item["table_count"] += table_count - item["total_chars"] += total_chars - - await _attach_connected_asset_counts( - db, - document_id=document_id, - job_result_id=job_result_id, - items_by_path=items_by_path, - sid_to_path=sid_to_path, - ) - - # Root is a virtual navigation container. Media availability for the - # whole document is exposed through global SEARCH actions, not as Root - # node-local images/tables. - root_item = items_by_path.get("Root") - if root_item: - root_item["image_count"] = 0 - root_item["table_count"] = 0 - - -async def _load_direct_chunk_counts( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - all_section_ids: list[str], -) -> dict[str, tuple[int, int, int, int]]: - chunk_stmt = ( - select( - DocumentChunk.section_id, - func.count( - case( - (DocumentChunk.chunk_type.notin_(["image", "table"]), literal_column("1")), - ) - ).label("text_count"), - func.count( - case( - (DocumentChunk.chunk_type == "image", literal_column("1")), - ) - ).label("image_count"), - func.count( - case( - (DocumentChunk.chunk_type == "table", literal_column("1")), - ) - ).label("table_count"), - func.coalesce( - func.sum(func.length(DocumentChunk.content)), 0 - ).label("total_chars"), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_section_ids)) - .group_by(DocumentChunk.section_id) - ) - chunk_rows = (await db.execute(chunk_stmt)).all() - return { - section_id: (int(text_count), int(image_count), int(table_count), int(total_chars)) - for section_id, text_count, image_count, table_count, total_chars in chunk_rows - } - - -async def _attach_connected_asset_counts( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - items_by_path: dict[str, dict], - sid_to_path: dict[str, str], -) -> None: - scope_items_with_zero_assets = [ - item - for item in items_by_path.values() - if item["show_summary"] and item["image_count"] == 0 and item["table_count"] == 0 - ] - if not scope_items_with_zero_assets: - return - - scope_section_ids = { - item["section_id"] - for item in items_by_path.values() - if item.get("section_id") - } - if not scope_section_ids: - return - - connect_stmt = ( - select( - DocumentChunk.section_id, - DocumentChunk.chunk_metadata, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(list(scope_section_ids))) - .where(DocumentChunk.chunk_type == "text") - ) - connect_result = (await db.execute(connect_stmt)).all() - - section_target_ids: dict[str, set[str]] = {} - for section_id, metadata in connect_result: - if not isinstance(metadata, dict): - continue - for connection in metadata.get("connect_to") or []: - target_id = connection.get("target", "") - if target_id: - section_target_ids.setdefault(section_id, set()).add(target_id) - - if not section_target_ids: - return - - all_target_ids: set[str] = set() - for target_ids in section_target_ids.values(): - all_target_ids.update(target_ids) - - target_type_stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) - .where(DocumentChunk.chunk_type.in_(["image", "table"])) - ) - target_type_result = (await db.execute(target_type_stmt)).all() - target_types = {chunk_id: chunk_type for chunk_id, chunk_type in target_type_result} - - for section_id, target_ids in section_target_ids.items(): - ref_path = sid_to_path.get(section_id, "") - if not ref_path: - continue - referenced_images = sum(1 for target_id in target_ids if target_types.get(target_id) == "image") - referenced_tables = sum(1 for target_id in target_ids if target_types.get(target_id) == "table") - if referenced_images == 0 and referenced_tables == 0: - continue - for item_path, item in items_by_path.items(): - if not item["show_summary"]: - continue - if _chunk_belongs_to_item(ref_path, item_path): - item["image_count"] += referenced_images - item["table_count"] += referenced_tables - - -def _chunk_belongs_to_item(chunk_path: str, item_path: str) -> bool: - if item_path == "Root": - return chunk_path == item_path - return chunk_path == item_path or chunk_path.startswith(item_path + " / ") diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py deleted file mode 100644 index 83cb39fa7..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Prompt projection for agentic section navigation (Collector Agent model).""" -from __future__ import annotations - -from typing import Any - - -def format_nav_trace( - nav_trace: list[dict[str, Any]], -) -> str: - """Render the unified navigation trace block. - - Includes compact navigation history. Current collection state is rendered - separately by the Agent State block. - """ - if not nav_trace: - return "" - - lines = ["=== Navigation Trace ==="] - for entry in nav_trace: - step = entry.get("step", "?") - scope = entry.get("scope", "root") - action = entry.get("action", "?") - reason = entry.get("reason", "") - action_display = action - drill_into = entry.get("drill_into") - if action == "EXPAND" and drill_into: - action_display = f'EXPAND "{drill_into}"' - elif action == "BACK": - back_to = entry.get("back_to") - target = f'"{back_to}"' if back_to else "root" - action_display = f"BACK to {target}" - - lines.append(f"Step {step}: scope={scope} → {action_display}") - - # Show tool usage and results so LLM can avoid repeating searches - tool_results = entry.get("tool_results", {}) - if tool_results: - tool_name = tool_results.get("tool", "") - tool_query = tool_results.get("query", "") - matched = int(tool_results.get("matched") or 0) - tool_status = str(tool_results.get("status") or "") - status = ( - f"found {matched} match(es)" - if matched - else f"no matches ({tool_status})" - if tool_status - else "no matches" - ) - lines.append(f' 🔧 {tool_name}("{tool_query}") → {status}') - - # Show what was collected in this step - step_collected = entry.get("collected", []) - if step_collected: - paths_display = ", ".join(f'"{c}"' for c in step_collected) - lines.append(f" collected: {paths_display}") - - result_status = entry.get("result_status") - if result_status and result_status != "ok": - lines.append(f" result_status: {result_status}") - - if reason: - lines.append(f" reason: {reason}") - lines.append("") - - lines.append("=== End Trace ===") - return "\n".join(lines) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py deleted file mode 100644 index 92592d4f5..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Section-tree loading and prompt projection for agentic navigation.""" -from __future__ import annotations - -from loguru import logger -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import DocumentSection -from shared.services.retrieval.agentic.navigation.section_counts import attach_section_counts -from shared.services.retrieval.search.lexical_text import normalize_section_path, split_section_path - - -async def load_document_section_rows( - db: AsyncSession, - document_id: str, - job_result_id: str, -) -> list: - stmt = ( - select( - DocumentSection.section_id, - DocumentSection.section_title, - DocumentSection.section_path, - DocumentSection.summary, - DocumentSection.sort_order, - ) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .order_by(DocumentSection.sort_order) - ) - return list((await db.execute(stmt)).all()) - - -async def load_child_sections( - db: AsyncSession, - document_id: str, - job_result_id: str, - scope_path: str | list[str] | None = None, - exclude_paths: set[str] | None = None, - limit_depth: bool = True, - section_rows: list | None = None, -) -> list[dict]: - """Load the continuous context tree for a navigation scope.""" - if section_rows is None: - section_rows = await load_document_section_rows( - db, - document_id=document_id, - job_result_id=job_result_id, - ) - if not section_rows: - return [] - - if isinstance(scope_path, list): - scope_list = [normalize_section_path(path) for path in scope_path] - elif scope_path: - scope_list = [normalize_section_path(scope_path)] - else: - scope_list = [] - - scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0 - excluded_paths = exclude_paths or set() - - logger.debug( - f" load_child_sections: scopes={scope_list or ['root']} " - f"scope_depth={scope_depth} exclude_paths={excluded_paths if excluded_paths else 'none'} " - f"total_sections={len(section_rows)}" - ) - - all_sections: dict[str, dict] = {} - for section_id, title, path, summary, sort_order in section_rows: - if not path: - continue - normalized_path = normalize_section_path(path) - parts = split_section_path(normalized_path) - # Treat Root as a virtual L1 node so it appears in the - # navigation tree and the LLM can decide whether to drill in. - if normalized_path == "Root" and not parts: - parts = ["Root"] - all_sections[normalized_path] = { - "title": title or parts[-1] if parts else normalized_path, - "summary": summary or "", - "sort_order": int(sort_order or 0), - "section_id": section_id, - "parts": parts, - "depth": len(parts), - } - - ancestor_prefixes: set[str] = set() - for scope in scope_list: - scope_parts = split_section_path(scope) - for index in range(1, len(scope_parts) + 1): - ancestor_prefixes.add(" / ".join(scope_parts[:index])) - - items_by_path = _select_scope_items( - all_sections, - scope_list=scope_list, - ancestor_prefixes=ancestor_prefixes, - exclude_paths=excluded_paths, - ) - if not items_by_path: - return [] - - if limit_depth: - allowed_set = _resolve_allowed_depths(items_by_path, scope_list) - if allowed_set: - to_remove = [ - path - for path, item in items_by_path.items() - if item["show_summary"] and item["level"] not in allowed_set - ] - for path in to_remove: - del items_by_path[path] - - if not items_by_path: - return [] - - await attach_section_counts( - db, - document_id=document_id, - job_result_id=job_result_id, - all_sections=all_sections, - items_by_path=items_by_path, - ) - - sorted_items = sorted(items_by_path.values(), key=lambda item: item["sort_order"]) - for item in sorted_items: - item.pop("sort_order", None) - item.pop("section_id", None) - - # Mark leaf status (no descendants in full section list) - all_section_paths = set(all_sections.keys()) - for item in sorted_items: - item_path = item["path"] - item["is_leaf"] = not any( - path != item_path and path.startswith(item_path + " / ") - for path in all_section_paths - ) - return sorted_items - - -def _select_scope_items( - all_sections: dict[str, dict], - *, - scope_list: list[str], - ancestor_prefixes: set[str], - exclude_paths: set[str], -) -> dict[str, dict]: - items_by_path: dict[str, dict] = {} - - def is_excluded(path: str) -> bool: - return bool( - exclude_paths - and any(path == excluded or path.startswith(excluded + " / ") for excluded in exclude_paths) - ) - - for path, meta in all_sections.items(): - parts = meta["parts"] - depth = meta["depth"] - - if not scope_list: - if depth < 1 or is_excluded(path): - continue - items_by_path[path] = _make_item(path, meta, show_summary=True) - continue - - matched_scope = _find_matched_scope(parts, depth=depth, scope_list=scope_list) - if matched_scope: - if is_excluded(path): - continue - items_by_path[path] = _make_item(path, meta, show_summary=True) - continue - - max_scope_depth = max(len(split_section_path(scope)) for scope in scope_list) - if depth <= max_scope_depth: - if depth == 1 and path in ancestor_prefixes: - items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) - elif depth > 1: - parent_prefix = " / ".join(parts[:-1]) - if parent_prefix in ancestor_prefixes: - items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) - - return items_by_path - - -def _make_item(path: str, meta: dict, show_summary: bool) -> dict: - return { - "path": path, - "title": meta["title"], - "summary": meta["summary"], - "level": meta["depth"], - "sort_order": meta["sort_order"], - "chunk_count": 0, - "image_count": 0, - "table_count": 0, - "total_chars": 0, - "section_id": meta["section_id"], - "show_summary": show_summary, - } - - -def _find_matched_scope(parts: list[str], *, depth: int, scope_list: list[str]) -> str | None: - for scope in scope_list: - scope_parts = split_section_path(scope) - scope_depth = len(scope_parts) - if depth > scope_depth and parts[:scope_depth] == scope_parts: - return scope - return None - - -def _resolve_allowed_depths(items_by_path: dict[str, dict], scope_list: list[str]) -> set[int]: - if not scope_list: - depths = { - item["level"] - for item in items_by_path.values() - if item.get("show_summary", True) - } - return set(sorted(depths)[:2]) - - allowed_set: set[int] = set() - for scope in scope_list: - scope_parts = split_section_path(scope) - scope_depth = len(scope_parts) - child_depths = { - item["level"] - for item in items_by_path.values() - if item.get("show_summary", True) - and item["level"] > scope_depth - and split_section_path(item["path"])[:scope_depth] == scope_parts - } - if child_depths: - allowed_set.update(sorted(child_depths)[:2]) - return allowed_set diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py deleted file mode 100644 index 45ef16939..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py +++ /dev/null @@ -1,137 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.agentic.core.types import DocTreeNode -from shared.services.retrieval.agentic.navigation import assets as asset_tools -from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows -from shared.services.retrieval.hydration.path import hydrate_paths_to_rows -from shared.services.retrieval.hydration.reference import hydrate_referenced_chunk_rows - - -async def hydrate_path_selections_into_node( - db: AsyncSession, - *, - node: DocTreeNode, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str, - job_result_id: str | None = None, -) -> None: - chunks = await hydrate_paths_to_rows( - db, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - if not chunks: - return - - chunks = await _append_connected_asset_targets(db, chunks) - resolved_job_result_id = job_result_id or _find_job_result_id(chunks) - if resolved_job_result_id: - await _attach_root_asset_owners( - db, - document_id=document_id, - job_result_id=resolved_job_result_id, - chunks=chunks, - ) - - add_chunks_to_node(node, chunks) - - -async def hydrate_chunk_refs_into_node( - db: AsyncSession, - *, - node: DocTreeNode, - refs: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str, - job_result_id: str | None = None, -) -> None: - chunks = await hydrate_referenced_chunk_rows( - db=db, - user_id=user_id, - namespace=namespace, - refs=refs, - ) - if not chunks: - return - - chunks = await _append_connected_asset_targets(db, chunks) - resolved_job_result_id = job_result_id or _find_job_result_id(chunks) - if resolved_job_result_id: - await _attach_root_asset_owners( - db, - document_id=document_id, - job_result_id=resolved_job_result_id, - chunks=chunks, - ) - - add_chunks_to_node(node, chunks) - - -async def _append_connected_asset_targets( - db: AsyncSession, chunks: list[dict[str, Any]] -) -> list[dict[str, Any]]: - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if not connected: - return chunks - - owner_map = asset_tools.build_connected_owner_map(chunks) - for chunk in connected: - if not chunk.get("owner_section_path"): - chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or "")) - return [*chunks, *connected] - - -async def _attach_root_asset_owners( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - chunks: list[dict[str, Any]], -) -> None: - root_map = await asset_tools.resolve_root_asset_owners( - db, - document_id=document_id, - job_result_id=job_result_id, - chunks=chunks, - ) - if not root_map: - return - - for chunk in chunks: - if chunk.get("owner_section_path"): - continue - chunk_id = str(chunk.get("chunk_id") or "") - if chunk_id in root_map: - chunk["owner_section_path"] = root_map[chunk_id] - - -def _find_job_result_id(chunks: list[dict[str, Any]]) -> str | None: - return next( - (str(chunk["job_result_id"]) for chunk in chunks if chunk.get("job_result_id")), - None, - ) - - -def add_chunks_to_node(node: DocTreeNode, chunks: list[dict[str, Any]]) -> None: - for chunk in chunks: - real_path = ( - chunk.get("owner_section_path") - or chunk.get("section_path") - or chunk.get("source_chunk_path") - ) - if real_path: - node.add_leaf_chunks(str(real_path), [chunk]) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py deleted file mode 100644 index 83b6ac57f..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Per-document navigation state for the collector runtime.""" -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Literal - -from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger - -RejectReason = Literal["tool_adjudicated", "navigational_abandon"] - -# Strength ordering: a stronger reason overrides a weaker one so we keep the -# most informative record for a given path. -_REASON_STRENGTH: dict[RejectReason, int] = { - "tool_adjudicated": 2, - "navigational_abandon": 1, -} - - -@dataclass -class RejectionRecord: - """A single rejection entry in the navigation state ledger. - - Two reasons are distinguished: - - - ``tool_adjudicated``: a SEARCH_* tool reconciliation proved the path has - no matching assets. Strong (content-level) negative signal. Not revived - this round; future strong-signal revival is a TODO. - - ``navigational_abandon``: BACK left an unproductive scope. Weak (soft) - signal; any discovery / lexical hit revives the path. - """ - - path: str - reason: RejectReason - step: int - detail: str = "" - - -@dataclass -class NavigationState: - """Mutable state for one document navigation loop. - - Path state is tracked in two orthogonal dimensions: - - - **Coverage** (positive "already taken as evidence"): derived from - ``collected_paths`` via :meth:`covered_paths` / :meth:`outline_paths`. - - **Rejection** (negative "evaluated, not taken"): a single labelled - ledger in :attr:`rejected` keyed by normalized path. Replaces the - former ``rejected_paths`` / ``rejected_collect_paths`` dual sets. - """ - - document_id: str - document_name: str - job_result_id: str - current_scope: str | None = None - expanded_scopes: set[str] = field(default_factory=set) - rejected: dict[str, RejectionRecord] = field(default_factory=dict) - collected_paths: list[dict[str, Any]] = field(default_factory=list) - nav_trace: list[dict[str, Any]] = field(default_factory=list) - tool_history: list[dict[str, Any]] = field(default_factory=list) - blocked_asset_searches: set[str] = field(default_factory=set) - step_count: int = 0 - - # ── Coverage helpers (single source: collected_paths) ──────────────── - - def covered_paths(self) -> set[str]: - """Full-evidence paths (hydrate_mode != 'outline').""" - return { - PathLedger.normalize(str(item.get("path") or "")) - for item in self.collected_paths - if item.get("path") and item.get("hydrate_mode") != "outline" - } - - def outline_paths(self) -> set[str]: - """Outline-only paths; excludes any path also collected as full.""" - full = self.covered_paths() - return { - PathLedger.normalize(str(item.get("path") or "")) - for item in self.collected_paths - if item.get("path") - and item.get("hydrate_mode") == "outline" - and PathLedger.normalize(str(item.get("path") or "")) not in full - } - - # ── Snapshot / delta for replayable traces ────────────────────────── - - def snapshot_delta( - self, - *, - before_scope: str | None, - expanded_before: set[str], - rejected_before: dict[str, RejectionRecord], - collected_before_count: int, - ) -> dict[str, Any]: - rejected_added: list[dict[str, Any]] = [] - for path, record in self.rejected.items(): - if path in rejected_before: - continue - rejected_added.append({ - "path": record.path, - "reason": record.reason, - "step": record.step, - "detail": record.detail, - }) - rejected_added.sort(key=lambda item: (item["step"], item["path"])) - return { - "current_scope_before": before_scope or "root", - "current_scope_after": self.current_scope or "root", - "expanded_added": sorted(self.expanded_scopes - expanded_before), - "rejected_added": rejected_added, - "collected_added": [ - item.get("path", "") - for item in self.collected_paths[collected_before_count:] - if item.get("path") - ], - } - - # ── Mutation helpers ──────────────────────────────────────────────── - - def add_collected( - self, - item: dict[str, Any], - *, - step: int, - scope_context: str | None, - ) -> dict[str, Any]: - enriched = dict(item) - enriched["collected_at_step"] = step - enriched["scope_context"] = scope_context or "root" - self.collected_paths.append(enriched) - return enriched - - def mark_expanded(self, path: str | None) -> None: - normalized = PathLedger.normalize(path) - if normalized: - self.expanded_scopes.add(normalized) - - def mark_rejected_collect( - self, - path: str | None, - *, - step: int, - detail: str = "", - ) -> None: - """Record a tool-adjudicated rejection (strong, content-level).""" - normalized = PathLedger.normalize(path) - if not normalized: - return - self._upsert_rejection( - normalized, - reason="tool_adjudicated", - step=step, - detail=detail, - ) - - def mark_rejected_if_unproductive( - self, - path: str | None, - *, - step: int, - detail: str = "", - ) -> None: - """Record a soft navigational abandon when leaving an unproductive scope. - - Only written when no stronger record exists for the path. - """ - normalized = PathLedger.normalize(path) - if not normalized: - return - has_full_collect = any( - item.get("hydrate_mode") != "outline" - and PathLedger.is_same_or_descendant(item.get("path"), normalized) - for item in self.collected_paths - ) - if has_full_collect: - return - self._upsert_rejection( - normalized, - reason="navigational_abandon", - step=step, - detail=detail, - ) - - def _upsert_rejection( - self, - normalized_path: str, - *, - reason: RejectReason, - step: int, - detail: str, - ) -> None: - existing = self.rejected.get(normalized_path) - if existing is not None and _REASON_STRENGTH[existing.reason] >= _REASON_STRENGTH[reason]: - # Keep the stronger prior record. - return - self.rejected[normalized_path] = RejectionRecord( - path=normalized_path, - reason=reason, - step=step, - detail=detail, - ) - - def rejected_paths_with_reason(self, reason: RejectReason) -> set[str]: - """All paths rejected with a specific reason label.""" - return { - path for path, record in self.rejected.items() - if record.reason == reason - } - - def blocked_asset_types_for_scope(self, scope: str | None) -> set[str]: - prefix = f"{PathLedger.normalize(scope) or 'root'}:" - return { - key.split(":", 1)[1] - for key in self.blocked_asset_searches - if key.startswith(prefix) - } - - def block_asset_search(self, scope: str | None, asset_type: str) -> None: - normalized_scope = PathLedger.normalize(scope) or "root" - normalized_type = asset_type.strip().lower() - if normalized_type: - self.blocked_asset_searches.add(f"{normalized_scope}:{normalized_type}") diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py deleted file mode 100644 index 066fa1302..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py +++ /dev/null @@ -1,364 +0,0 @@ -"""Agentic retrieval navigation tools — observe-act collector model.""" -from __future__ import annotations - -from typing import Any - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.agentic.navigation.assets import ( - count_assets_under_scope, -) -from shared.services.retrieval.agentic.core.budget import ( - BudgetExceeded, - budget_status_from_snapshot, -) -from shared.services.retrieval.agentic.prompts import ( - COLLECTOR_PROMPT, - adjust_budget_snapshot, - parse_collector_response, -) -from shared.services.retrieval.agentic.navigation.section_prompt_projection import ( - format_nav_trace, -) -from shared.services.retrieval.agentic.navigation.actions import ( - build_legal_actions, - format_actionable_observation, - format_agent_state_block, -) -from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections -from shared.services.retrieval.agentic.navigation.state import RejectionRecord -from shared.services.retrieval.agentic.core.types import DocTreeNode, NavigateStepResult -from shared.services.retrieval.llm_adapter import LLMFn - - - -async def navigate_step( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - query: str, - llm_fn: LLMFn, - user_id: str, - namespace: str, - doc_name: str = "", - scope_path: str | None = None, - exclude_paths: set[str] | None = None, - budget_snapshot: dict | None = None, - nav_trace: list[dict[str, Any]] | None = None, - collected_paths: list[dict[str, Any]] | None = None, - expanded_scopes: set[str] | None = None, - rejected: dict[str, RejectionRecord] | None = None, - disabled_asset_types: set[str] | None = None, - discovery_hints: list[dict[str, Any]] | None = None, - section_rows: list | None = None, - query_intent: str = "UNKNOWN", - search_context: str = "", - prior_tool_result: dict[str, Any] | None = None, -) -> NavigateStepResult: - """Navigate one document scope with a single observe-act decision.""" - scope_paths = [scope_path] if scope_path else [] - - try: - items = await load_child_sections( - db, - document_id, - job_result_id, - scope_path, - exclude_paths=exclude_paths, - section_rows=section_rows, - ) - if not items: - return NavigateStepResult.stop( - scope_paths[0] if scope_paths else None, - reason="No visible sections in the current scope.", - ) - - budget_status = budget_status_from_snapshot(budget_snapshot) - if budget_status in {"CRITICAL", "EXHAUSTED"}: - total_images, total_tables = 0, 0 - else: - total_images, total_tables = await count_assets_under_scope( - db, - document_id=document_id, - job_result_id=job_result_id, - scope_paths=scope_paths, - ) - - expanded_path_set = set(expanded_scopes or _expanded_paths_from_trace(nav_trace or [])) - if scope_path: - expanded_path_set.add(scope_path) - rejection_ledger = rejected or {} - provisional_action_set = build_legal_actions( - items=items, - current_scope=scope_path, - collected_paths=collected_paths or [], - expanded_scopes=expanded_path_set, - discovery_hints=discovery_hints, - rejected=rejection_ledger, - total_images=total_images, - total_tables=total_tables, - disabled_asset_types=disabled_asset_types or set(), - budget_snapshot=budget_snapshot, - ) - provisional_observation_text, provisional_overflowed = ( - format_actionable_observation( - items=items, - action_set=provisional_action_set, - ) - ) - - trace_block = format_nav_trace(nav_trace or []) - - # Estimate this call's prompt token cost and adjust the budget - # snapshot so the LLM sees post-call budget, not pre-call. - # This prevents the LLM from seeing misleadingly low percentages - # (e.g. 63% when it will actually be 89% after this call). - prompt_tokens_est = ( - len(provisional_observation_text) - + len(trace_block) - + 800 - ) // 2 # rough chars-to-tokens ratio - adjusted_snapshot = adjust_budget_snapshot( - budget_snapshot, prompt_tokens_est, - ) - if ( - budget_status_from_snapshot(adjusted_snapshot) - == budget_status_from_snapshot(budget_snapshot) - ): - action_set = provisional_action_set - actionable_observation = provisional_observation_text - overflowed = provisional_overflowed - else: - action_set = build_legal_actions( - items=items, - current_scope=scope_path, - collected_paths=collected_paths or [], - expanded_scopes=expanded_path_set, - discovery_hints=discovery_hints, - rejected=rejection_ledger, - total_images=total_images, - total_tables=total_tables, - disabled_asset_types=disabled_asset_types or set(), - budget_snapshot=adjusted_snapshot, - ) - actionable_observation, overflowed = format_actionable_observation( - items=items, - action_set=action_set, - ) - observation = { - "visible_sections": [ - item.get("path", "") - for item in items - if item.get("path") - ][:50], - "available_images": total_images, - "available_tables": total_tables, - "prior_tool_result": prior_tool_result, - "current_scope": scope_path or "root", - "query_intent": query_intent, - "legal_actions": { - "expand": [item.id for item in action_set.expand], - "collect": [item.id for item in action_set.collect], - "back": [item.id for item in action_set.back], - "search": [item.id for item in action_set.search], - "finish": [action_set.finish.id] if action_set.finish else [], - }, - "rejected": { - path: {"reason": record.reason, "step": record.step} - for path, record in rejection_ledger.items() - }, - } - agent_state_block = format_agent_state_block( - current_scope=scope_path, - query_intent=query_intent, - expanded_scopes=expanded_path_set, - rejected=rejection_ledger, - collected_paths=collected_paths or [], - prior_tool_result=prior_tool_result, - search_context=search_context, - budget_snapshot=adjusted_snapshot, - ) - - prompt = COLLECTOR_PROMPT.format( - doc_name=doc_name or document_id, - doc_id=document_id, - agent_state_block=agent_state_block, - trace_block=trace_block, - query=query, - actionable_observation=actionable_observation, - ) - - response = await llm_fn(prompt) - parsed = parse_collector_response(response) - requested_action = parsed["action"] - selected_tools = parsed["tools"] - tool_params = parsed.get("tool_params", {}) - reason = parsed.get("reason", "") - raw_collect = parsed.get("collect", []) - action_id = parsed.get("action_id") - legal_main = action_set.get(action_id) - action = ( - legal_main.action - if legal_main and legal_main.action != "COLLECT" - else requested_action - ) - if action != requested_action: - reason = ( - f"Action field '{requested_action}' did not match legal action " - f"ID '{action_id}'; executing ID-defined action '{action}'. " - + reason - ).strip()[:500] - selected_tools = [action] if action in ("SEARCH_IMAGES", "SEARCH_TABLES") else [] - - scope_label = scope_path or "root" - logger.info( - f" navigate_step scope={scope_label}: " - f"action={action} collect={len(raw_collect)} " - f"action_id={action_id} tools={selected_tools} " - f"tool_params={tool_params} " - f"overflowed={overflowed}" - ) - - node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None) - node.outline_items = [item for item in items if item.get("show_summary", True)] - - # Resolve COLLECT side effects from legal action IDs. - valid_collect: list[dict[str, Any]] = [] - invalid_collect: list[str] = [] - existing_collect_modes = _existing_collect_modes(collected_paths or []) - for item in raw_collect: - collect_id = item.get("id") - legal_collect = action_set.get(collect_id) - if legal_collect and legal_collect.action == "COLLECT" and legal_collect.path: - confidence = item.get("confidence", 0.7) - outline = bool(item.get("outline", False)) and ( - budget_status_from_snapshot(adjusted_snapshot) != "CRITICAL" - or query_intent in {"MACRO_SUMMARY", "STRUCTURE_OVERVIEW"} - ) - if ( - legal_main - and legal_main.action == "EXPAND" - and legal_main.path == legal_collect.path - ): - outline = True - hydrate_mode = "outline" if outline else "chunks" - existing_mode = existing_collect_modes.get(legal_collect.path) - if existing_mode == "chunks": - continue - if existing_mode == "outline" and hydrate_mode == "outline": - continue - node.confidence[legal_collect.path] = confidence - valid_collect.append({ - "path": legal_collect.path, - "confidence": confidence, - "hydrate_mode": hydrate_mode, - }) - elif collect_id: - invalid_collect.append(str(collect_id)) - - valid_drill: list[dict[str, Any]] = [] - result_status = "ok" - result_note: str | None = None - drill_into: str | None = None - back_to: str | None = None - if requested_action == "ERROR": - result_status = "invalid_response" - result_note = reason or "invalid model response" - elif action == "EXPAND": - if legal_main and legal_main.action == "EXPAND" and legal_main.path: - drill_into = legal_main.path - valid_drill.append({ - "path": drill_into, - "confidence": 0.8, - }) - else: - result_status = "invalid_action_id" - result_note = f"invalid_expand_id: {action_id}" - elif action == "BACK": - if legal_main and legal_main.action == "BACK": - back_to = legal_main.target_scope - else: - result_status = "invalid_action_id" - result_note = f"invalid_back_id: {action_id}" - elif action in ("SEARCH_IMAGES", "SEARCH_TABLES"): - if legal_main is None or legal_main.action != action: - result_status = "invalid_action_id" - result_note = f"invalid_search_id: {action_id}" - elif action == "FINISH": - if legal_main is None or legal_main.action != "FINISH": - result_status = "invalid_action_id" - result_note = f"invalid_finish_id: {action_id}" - - if invalid_collect and result_status == "ok": - result_status = "invalid_collect" - result_note = "invalid_collect_ids: " + ", ".join(invalid_collect[:5]) - - # Parse tool parameters for SEARCH - search_assets_params: dict[str, Any] | None = None - - if action in ("SEARCH_IMAGES", "SEARCH_TABLES") and result_status == "ok": - asset_type = "image" if action == "SEARCH_IMAGES" else "table" - collected_scope_paths = [ - str(item.get("path") or "") - for item in valid_collect - if item.get("path") - ] - search_assets_params = { - "query": query.strip(), - "asset_type": asset_type, - "scope_paths": collected_scope_paths or scope_paths, - } - elif action in ("SEARCH_IMAGES", "SEARCH_TABLES"): - selected_tools = [] - - return NavigateStepResult( - action=action, - collect=valid_collect, - drill=valid_drill, - back_to=back_to, - tools=selected_tools, - node=node, - reason=reason, - search_assets_params=search_assets_params, - observation=observation, - result_status=result_status, - result_note=result_note, - ) - - except BudgetExceeded: - raise - except Exception as exc: - logger.error(f" navigate_step failed for doc={document_id}: {exc}") - return NavigateStepResult.error( - scope_paths[0] if scope_paths else None, - reason=str(exc), - ) - - -def _expanded_paths_from_trace(nav_trace: list[dict[str, Any]]) -> set[str]: - expanded: set[str] = set() - for entry in nav_trace: - if entry.get("action") != "EXPAND": - continue - if entry.get("result_status", "ok") != "ok": - continue - drill_into = entry.get("drill_into") - if isinstance(drill_into, str) and drill_into: - expanded.add(drill_into) - return expanded - - -def _existing_collect_modes(collected_paths: list[dict[str, Any]]) -> dict[str, str]: - modes: dict[str, str] = {} - for item in collected_paths: - path = str(item.get("path") or "") - if not path: - continue - hydrate_mode = str(item.get("hydrate_mode") or "chunks") - if hydrate_mode != "outline": - modes[path] = "chunks" - elif modes.get(path) != "chunks": - modes[path] = "outline" - return modes diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py deleted file mode 100644 index 32a604966..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ /dev/null @@ -1,448 +0,0 @@ -"""Retrieval Agent orchestrator — evidence-only navigation loop. - -Flow: - Phase 1: Document selection (bottom_discovery + kg_document_select) - Phase 2: Per-document navigation (iterative BFS via navigate_step) - Phase 3: Render evidence text for downstream agents - -The orchestrator drives navigation through a per-document observe-act loop. -Each navigate_step is a single LLM call that chooses one main action -(EXPAND/BACK/SEARCH_IMAGES/SEARCH_TABLES/FINISH) plus optional collection -side effects. FINISH explicitly terminates navigation for that document. - -KNOWHERE does not generate final answers. Downstream agents decide whether the -returned evidence is sufficient for their task and may call retrieval again. -""" -from __future__ import annotations - -import os -from typing import Any - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.services.retrieval.agentic.core.budget import BudgetLedger -from shared.services.retrieval.agentic.discovery.phase import ( - run_initial_discovery, -) -from shared.services.retrieval.agentic.navigation.document import DocumentNavigationRunner -from shared.services.retrieval.agentic.evidence.builder import ( - trim_evidence_to_budget as _trim_evidence_to_budget, -) -from shared.services.retrieval.agentic.core.runtime import ( - AgentLlmBudget, - build_config_from_env as _build_config_from_env, - load_budget_inventory as _load_budget_inventory, -) -from shared.services.retrieval.agentic.core.trace import TraceRecorder -from shared.services.retrieval.agentic.core.types import ( - AgentRunConfig, - AgentState, - AgenticResult, - DecisionTraceStep, -) -from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.settings import DEFAULT_TOP_K, resolve_disabled_asset_types - - -class RetrievalAgent: - """Agentic retrieval orchestrator — navigate and return evidence. - - Usage:: - - agent = RetrievalAgent() - result = await agent.run( - db, user_id=..., namespace=..., query=..., llm_fn=..., ... - ) - # result.evidence_text — hierarchical context for downstream agents - # result.answer_text — deprecated, always empty - # result.referenced_chunks — chunk IDs for hit stats / frontend - - The agent requires a valid ``llm_fn`` for LLM-driven navigation. - If ``llm_fn`` is None, the run returns discovery-only results. - """ - - async def run( - self, - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int = DEFAULT_TOP_K, - llm_fn: LLMFn | None = None, - exclude_document_ids: list[str] | None = None, - exclude_sections: list[dict[str, str]] | None = None, - chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = 'delete', - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, - internal_recall_k: int | None = None, - config: AgentRunConfig | None = None, - ledger: BudgetLedger | None = None, - parent_run_id: str | None = None, - workflow_step_id: str | None = None, - ) -> AgenticResult: - """Run the agentic retrieval pipeline. - - Returns an ``AgenticResult`` containing rendered evidence text - and referenced chunk IDs. Never raises — errors are captured in - trace and the best available evidence is returned. - """ - config = config or _build_config_from_env() - exclude_document_ids = exclude_document_ids or [] - exclude_sections = exclude_sections or [] - - state = AgentState() - state.ledger = ledger or BudgetLedger( - total=config.token_budget_total, - planning_ratio=config.planning_ratio, - bootstrap=config.bootstrap_budget, - per_doc_min_share=config.per_doc_min_share, - ) - total_chunks, total_docs, chunks_count_by_doc = await _load_budget_inventory( - db, - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, - ) - state.kg_total_chunks = total_chunks - state.kg_total_docs = total_docs - state.ledger.total_chunks = total_chunks - state.ledger.total_docs = total_docs - trace = TraceRecorder( - db, user_id=user_id, namespace=namespace, query=query, - config=config, top_k=top_k, chunk_types=chunk_types, - filters={ - 'exclude_document_ids': exclude_document_ids, - 'exclude_sections': exclude_sections, - 'signal_paths': signal_paths, - 'internal_recall_k': internal_recall_k, - }, - parent_run_id=parent_run_id, - workflow_step_id=workflow_step_id, - ) - - trace_enabled = os.environ.get('RETRIEVAL_AGENTIC_TRACE_ENABLED', 'true') == 'true' - if trace_enabled: - await trace.create_run() - - logger.info( - f'agentic retrieval START: query="{query[:60]}..." ' - f'top_k={top_k} latency_budget={config.latency_budget_ms}ms ' - f'token_budget={config.token_budget_total}' - ) - - if llm_fn is None: - logger.warning('agentic: no llm_fn provided — running discovery-only mode') - - bootstrap_llm_fn: LLMFn | None = None - llm_budget = AgentLlmBudget(state) - if llm_fn is not None: - bootstrap_llm_fn = llm_budget.for_pool(llm_fn, pool='bootstrap') - - discovery_rows = await run_initial_discovery( - db, - state=state, - trace=trace, - trace_enabled=trace_enabled, - user_id=user_id, - namespace=namespace, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - chunk_types=chunk_types, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, - internal_recall_k=internal_recall_k, - bootstrap_llm_fn=bootstrap_llm_fn, - ) - - # If no LLM or no docs selected, return discovery rows directly - if not state.selected_docs: - logger.info('agentic: no documents selected — returning discovery results') - no_docs_trace: list[dict[str, Any]] = [] - kg_select_step = DecisionTraceStep( - step_index=0, - agent='doc_selector', - parent_step_index=None, - phase='kg_select', - document_id=None, - document=None, - scope='corpus', - observation={ - 'query': query, - 'candidate_documents': len(state.doc_id_to_name), - }, - decision={ - 'action': 'select_documents', - 'args': {}, - 'reason': 'No documents selected for navigation', - }, - result={ - 'status': 'empty', - 'collected': [], - }, - budget=state.ledger.snapshot() if state.ledger else None, - elapsed_ms=state.elapsed_ms, - ) - no_docs_trace.append(kg_select_step.to_dict()) - if trace_enabled: - trace.record_decision_trace_step(kg_select_step) - discovery_refs = [ - { - 'chunk_id': r.get('chunk_id', ''), - 'document_id': r.get('document_id', ''), - 'chunk_type': r.get('chunk_type', ''), - 'section_path': ( - r.get('source_file_name', '') - if r.get('section_path') == 'Root' - else r.get('section_path', '') - ), - 'file_path': r.get('file_path', ''), - 'job_id': r.get('job_id', ''), - } - for r in discovery_rows[:top_k] - if r.get('chunk_id') - ] - terminal_step = DecisionTraceStep( - step_index=len(no_docs_trace), - agent='retrieval_agent', - parent_step_index=None, - phase='terminal', - document_id=None, - document=None, - scope='retrieve_step', - observation={ - 'router_used': 'agentic_discovery_only', - 'referenced_chunks': len(discovery_refs), - 'evidence_chars': 0, - }, - decision={ - 'action': 'complete', - 'args': {}, - 'reason': 'no_documents_selected', - }, - result={ - 'status': 'ok', - 'stop_reason': 'no_documents_selected', - 'failure_reason': '', - }, - budget=state.ledger.snapshot() if state.ledger else None, - elapsed_ms=state.elapsed_ms, - ) - no_docs_trace.append(terminal_step.to_dict()) - if trace_enabled: - trace.record_decision_trace_step(terminal_step) - if trace_enabled: - await trace.complete( - discovery_rows, - 'agentic_discovery_only', - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - return AgenticResult( - evidence_text='', - answer_text='', - referenced_chunks=discovery_refs, - router_used='agentic_discovery_only', - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - stop_reason='no_documents_selected', - decision_trace=no_docs_trace, - ) - - discovery_by_doc: dict[str, list[dict[str, Any]]] = {} - for row in discovery_rows: - doc_id = row.get('document_id', '') - # Root chunks are navigable via the section tree; exclude them - # from discovery hints where the bare "Root" label gives the LLM - # no actionable information. - section_path = str(row.get('section_path', '') or '').strip() - if not doc_id or section_path == 'Root': - continue - discovery_by_doc.setdefault(doc_id, []).append(row) - - - if state.ledger is not None: - await state.ledger.allocate_doc_caps({ - doc.document_id: chunks_count_by_doc.get(doc.document_id, 1) - for doc in state.selected_docs - }) - - # Phase 2 + 3: navigate once, then render evidence for downstream agents. - evidence_text = '' - stop_reason = 'evidence_only' - failure_reason = '' - decision_trace: list[dict[str, Any]] = [] - - # Record KG document selection as the first decision trace entry - if state.selected_docs: - kg_select_step = DecisionTraceStep( - step_index=0, - agent='doc_selector', - parent_step_index=None, - phase='kg_select', - document_id=None, - document=None, - scope='corpus', - observation={ - 'query': query, - 'candidate_documents': len(state.doc_id_to_name), - }, - decision={ - 'action': 'select_documents', - 'args': {}, - 'reason': f'{len(state.selected_docs)} document(s) selected for navigation', - }, - result={ - 'status': 'ok', - 'collected': [ - { - 'document': doc.source_file_name, - 'document_id': doc.document_id, - 'confidence': doc.confidence, - 'reason': doc.reason, - 'source': doc.source, - } - for doc in state.selected_docs - ], - }, - budget=state.ledger.snapshot() if state.ledger else None, - elapsed_ms=state.elapsed_ms, - ) - decision_trace.append(kg_select_step.to_dict()) - if trace_enabled: - trace.record_decision_trace_step(kg_select_step) - - if state.elapsed_ms >= config.latency_budget_ms: - stop_reason = 'latency_budget' - else: - navigation_runner = DocumentNavigationRunner( - db=db, - state=state, - trace=trace, - trace_enabled=trace_enabled, - user_id=user_id, - namespace=namespace, - query=query, - config=config, - discovery_by_doc=discovery_by_doc, - llm_fn=llm_fn, - llm_budget=llm_budget, - disabled_asset_types=resolve_disabled_asset_types(chunk_types), - ) - await navigation_runner.navigate_selected_documents() - decision_trace.extend( - _offset_decision_trace( - navigation_runner.decision_steps, - offset=len(decision_trace), - ) - ) - - context_remaining = state.ledger.remaining('context') if state.ledger else config.token_budget_total - evidence_text = await _trim_evidence_to_budget( - db, - doc_trees=state.doc_trees, - doc_id_to_name=state.doc_id_to_name, - context_remaining=context_remaining, - user_id=user_id, - namespace=namespace, - ledger=state.ledger, - ) - - # ══════════════════════════════════════════════════════════════════ - # Final Assembly - # ══════════════════════════════════════════════════════════════════ - router_used = ( - 'agentic_llm' if any(t.has_content() for t in state.doc_trees.values()) - else 'agentic_discovery_only' - ) - - # Collect referenced chunk IDs from all doc trees - all_refs: list[dict[str, Any]] = [] - seen_ref_ids: set[str] = set() - for doc_id, doc_tree in state.doc_trees.items(): - doc_name = state.doc_id_to_name.get(doc_id, doc_id) - for ref in doc_tree.collect_referenced_ids(document_name=doc_name): - cid = ref.get('chunk_id', '') - if cid and cid not in seen_ref_ids: - seen_ref_ids.add(cid) - all_refs.append(ref) - - terminal_step = DecisionTraceStep( - step_index=len(decision_trace), - agent='retrieval_agent', - parent_step_index=None, - phase='terminal', - document_id=None, - document=None, - scope='retrieve_step', - observation={ - 'router_used': router_used, - 'referenced_chunks': len(all_refs), - 'evidence_chars': len(evidence_text), - }, - decision={ - 'action': 'complete', - 'args': {}, - 'reason': stop_reason or failure_reason or 'retrieval_complete', - }, - result={ - 'status': 'error' if failure_reason else 'ok', - 'stop_reason': stop_reason, - 'failure_reason': failure_reason, - }, - budget=state.ledger.snapshot() if state.ledger else None, - elapsed_ms=state.elapsed_ms, - ) - decision_trace.append(terminal_step.to_dict()) - if trace_enabled: - trace.record_decision_trace_step(terminal_step) - - result = AgenticResult( - evidence_text=evidence_text, - answer_text='', - referenced_chunks=all_refs, - router_used=router_used, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - stop_reason=stop_reason, - failure_reason=failure_reason, - decision_trace=decision_trace, - ) - - logger.info( - f'agentic retrieval DONE: {len(all_refs)} referenced chunks, ' - f'evidence_text={len(evidence_text)} chars, ' - f'router={router_used}, steps={state.step_count}, ' - f'stop_reason={stop_reason}, ' - f'{state.elapsed_ms}ms' - ) - - if trace_enabled: - await trace.complete( - all_refs, - router_used, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - - return result - - -def _offset_decision_trace( - steps: list[dict[str, Any]], - *, - offset: int, -) -> list[dict[str, Any]]: - adjusted: list[dict[str, Any]] = [] - for step in steps: - copied = dict(step) - old_index = int(copied.get('step_index') or 0) - copied['step_index'] = old_index + offset - parent = copied.get('parent_step_index') - if parent is not None: - copied['parent_step_index'] = int(parent) + offset - adjusted.append(copied) - return adjusted diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py deleted file mode 100644 index 3a7c5b2cc..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/prompts.py +++ /dev/null @@ -1,316 +0,0 @@ -"""Prompt templates and response parsers for agentic retrieval.""" -from __future__ import annotations - -import json -import re -from typing import Any - -from shared.services.retrieval.agentic.core.budget import project_budget_snapshot - - -FILE_SELECT_PROMPT = """\ -You are a document routing assistant. - -{budget_block} -Below is a document corpus overview showing all available documents, -their navigation summaries, chunk counts, and media counts. -Some documents may show "🔍 Discovery hints" — these are preliminary keyword -matches from bottom-up search. Consider them as additional signals but make -your own judgment on document relevance. - -=== Document Corpus Overview === -{overview} -=== End Overview === - -User query: {query} -Based on the query, select documents that may contain relevant information. -If NO document in the corpus is relevant to the query, return an EMPTY array []. -Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"] -Do not include any explanation. -""" - - -COLLECTOR_PROMPT = """\ -You are a document navigation agent running an observe-act loop. - -Document: "{doc_name}" (id: {doc_id}) - -{agent_state_block} - -{trace_block} - -User query: {query} - -{actionable_observation} - -=== Rules === - -Each step chooses exactly ONE main action, plus optional COLLECT side effects. - -Action semantics: - - EXPAND observes a listed section's children in the next step. - - COLLECT adds a listed section and all descendant content to evidence. - - BACK only changes current scope; it does not collect evidence. - - SEARCH_IMAGES and SEARCH_TABLES inspect assets in the current scope. - Use only the listed SEARCH action ID. The asset inspector receives the - user's original query directly. - After a SEARCH result returns matches, use the matched assets and owner - sections to decide whether more owner context is needed; avoid repeating - the same asset search unless the current scope has changed and the prior - result is insufficient for the query. - - FINISH ends navigation for this document. - -COLLECT side effect: - - COLLECT includes the section AND ALL its descendant content. - - Set "outline": true to collect only structure (titles + summaries), - keeping children available for further EXPAND or COLLECT. - - If you COLLECT the same section you EXPAND as the main action, use - "outline": true so the section remains open for child exploration. - - If the advisory query intent is MACRO_SUMMARY or STRUCTURE_OVERVIEW - (document overview, chapter map, high-level summary), prefer outline - collection; outline evidence can be sufficient final evidence. - - If the advisory query intent is FACTUAL_DETAIL, NUMERIC_DETAIL, or - ASSET_LOOKUP, prefer full evidence collection ("outline": false), or - SEARCH_IMAGES/SEARCH_TABLES when visual/table evidence is central. - - If the advisory query intent is UNKNOWN, decide from the user's wording: - broad summaries can use outline, specific facts/numbers/assets need full - evidence. - - FINISH only when the collected evidence is sufficient for the user's query. - The system will not infer missing evidence for you. - - In CRITICAL budget mode, exploration is closed. Prefer the smallest - sufficient COLLECT side effects, then FINISH. - - In EXHAUSTED or overdraft budget mode, do not explore or search again. - Use current observations/tool results to FINISH, or collect only - indispensable visible evidence before FINISH. - - For [Leaf] nodes or small sections, prefer COLLECT over EXPAND. - -=== End Rules === - -Return ONLY a JSON object: -{{"collect": [{{"id": "C1", "confidence": , "outline": false}}], - "action": "", - "action_args": {{"id": "
"}}, - "reason": "..."}} -Do not include any explanation outside the JSON. - -IMPORTANT: -1. All agent-generated text (e.g., "reason" and other free-text fields) MUST be written in English. -2. Document content and section paths MUST remain in their original language. -3. Use only action IDs from Actionable Observation. Never invent IDs or write raw section paths as action targets. -4. The action value MUST match the chosen ID group: E*=EXPAND, B*=BACK, S*=SEARCH, F*=FINISH. -5. When Budget mode is CRITICAL or EXHAUSTED, choose the best sufficient COLLECT side effects and then FINISH. -""" - - -QUERY_INTENT_PROMPT = """\ -Classify the user's retrieval query for document navigation. - -Return ONLY a JSON object: {{"intent": "