diff --git a/AGENTS.md b/AGENTS.md index 503a16ad..fa13a28f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -103,13 +103,12 @@ flowchart TB subgraph RETRIEVE["⑤ Retrieval (shared)"] Query["POST /v1|/v2 retrieval/query"] --> Pipeline["run_retrieval_query"] Pipeline --> Classic["classic_topk / small_corpus (use_agentic=False)"] - Pipeline --> MapNav["mapnav checklist (default / use_agentic≠False)"] + Pipeline --> Explore["agent_explore + cursor_sdk (default / use_agentic≠False)"] Classic --> Channels["map_unit_discovery: path+content BM25 -> RRF"] Channels --> Rank["rank_retrieval_candidates"] - MapNav --> NavSnap["nav_snapshot + run_nav_episode"] - NavSnap --> Bridge["nav_bridge referenced_chunks"] + Explore --> Tools["corpus.* tools + harness.run_episode"] Rank --> Assemble["assemble_retrieval_results"] - Bridge --> Assemble + Tools --> Assemble Assemble --> Results["Cited Evidence Results"] end ``` @@ -493,8 +492,6 @@ debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting. | `content_search_text` | `Text` | Pre-tokenized for BM25 content channel | | `path_search_text` | `Text` | Pre-tokenized for BM25 path channel | | `term_search_text` | `Text` | Pre-tokenized for term/grep channel | -| `content_search_tsv` | `TSVECTOR` (computed) | PostgreSQL GIN index for full-text | -| `path_search_tsv` | `TSVECTOR` (computed) | PostgreSQL GIN index for path | | `source_chunk_path` | `Text` | Original parser path | | `file_path` | `Text` | Asset reference (`images/x.jpg`) | | `chunk_metadata` | `JSON` | Keywords, tokens, connect_to, etc. | @@ -545,23 +542,23 @@ debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting. Core retrieval internals are grouped by ownership: -- `execution/`: request shaping, route selection (classic / mapnav / small_corpus), and public response projection. -- `search/`: `map_unit_discovery` (persisted map-unit BM25 discovery, with a legacy chunk-level PG FTS fallback), scoring, section filters, candidate ranking. +- `execution/`: request shaping, route selection (classic / agent_explore / small_corpus), and public response projection. +- `search/`: `map_unit_discovery` (persisted map-unit BM25 discovery; incomplete or incompatible indexes raise), scoring, section filters, candidate ranking. - `hydration/`: row/path/reference hydration, inline assets, and result assembly. -- `nav/` + `nav_*.py`: map-nav checklist episode (PLANNER / HARVEST / CONTROL). +- `agent_explore/` + `agent_tools/`: default agentic route (cursor_sdk harness). Map-nav is archived under `deprecated/mapnav/`. - `trace/`: `DecisionTraceStep` mapping and `TraceRecorder`. - `graph/`: document graph publication/query support. - `stats/`: retrieval hit recording. ### Two Retrieval Modes -Per-request `use_agentic`: `False` → classic top-K (map-unit BM25); `None`/`True` → map-nav (default). +Per-request `use_agentic`: `False` → classic top-K (map-unit BM25); `None`/`True` → agent_explore (default harness `cursor_sdk`). -#### Classic Mode (map-unit BM25 + legacy FTS fallback) +#### Classic Mode (map-unit BM25) Primary path is `search.map_unit_discovery.map_unit_discovery`: Python BM25Okapi over the persisted `document_map_unit_tokens` index, path and content channels -only, fused by RRF. +only, fused by RRF. An incomplete or incompatible map-unit index raises. ```mermaid flowchart LR @@ -577,15 +574,15 @@ flowchart LR `score = weight / (k + rank + 1)` per channel, summed across channels, `k=60`. There is no scored term channel in this primary path. `term_search_text` / -`term_search_text_lower` are persisted at publish time but are only read by -the **legacy fallback** (`_legacy_chunk_discovery`), which runs only when a -revision's map-unit index is missing or incomplete: a single SQL query -scoring `GREATEST(ts_rank_cd(path_search_tsv), 2 * ts_rank_cd(content_search_tsv))` -OR `term_search_text LIKE '%query%'`, not three independently-ranked channels. +`term_search_text_lower` are persisted at publish time and are read by +`corpus.recall`'s term channel, not by classic discovery. -#### Map-nav Mode (default) +#### Agent-explore Mode (default) -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`. +Default agentic path is `agent_explore`: a `corpus.*` tool loop run by the +resolved harness (`AGENT_EXPLORE_HARNESS`, default `cursor_sdk`). Finish refs +are resolved to chunks, then assembled like classic. Map-nav is archived at +`deprecated/mapnav/` and is no longer a live route. ### Result Assembly diff --git a/README.md b/README.md index 8f0178ca..a9fdf7fc 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -20260506-102713 +Knowhere 2.0

Prepare unstructured data for AI Agents

@@ -37,7 +37,7 @@ It ingests unstructured documents and produces persistent, navigable memory: parsing, hierarchy reconstruction, multi-modal structuring, and graph construction in a single pipeline. Every result stays connected to its document, section, source pages, and related assets, making the output a natural fit for *Agentic RAG*, *vector-based RAG*, or any LLM workflow. -Knowhere supports complementary **Vision and Text tracks**. Text-native documents retain precise extracted structure, while complex PDFs and PowerPoint files can be understood directly as pages by frontier vision models. Both tracks converge into the same memory schema, hierarchy, retrieval engine, and citation model. +Knowhere 2.0 supports complementary **Vision and Text tracks**. Text-native documents retain precise extracted structure, while complex PDFs and PowerPoint files can be understood directly as pages by frontier vision models. Both tracks converge into the same memory schema, hierarchy, retrieval engine, and citation model. > [!NOTE] > **Get started in seconds with Knowhere Cloud.** @@ -45,6 +45,7 @@ Knowhere supports complementary **Vision and Text tracks**. Text-native document ## 📢 News +- **September 8, 2026**: 🧭 **Introducing agent-native Retrieval 2.0.** MapNav has evolved from a fixed navigation workflow into a corpus-native foundation for agents. Knowhere provides a unified schema, hierarchy-aware tools, and resolvable evidence references; the agent decides how to search, traverse, read, and cite. The same foundation powers built-in retrieval and external agents through MCP. - **September 2026**: 👁️ **Introducing dual-track Document Parsing 2.0.** Vision Page and Text Track now converge into one hierarchy-native memory schema for retrieval, understanding, and citation. - **June 1, 2026**: 📚 **Knowhere now supports ultra-long PDFs and atlas-style documents.** The parsing pipeline can process long-form PDFs with hundreds of pages (for example, 300, 500, or more) and route technical atlases or drawing collections through a dedicated layout-aware parser. - **May 7, 2026**: 🚀 **Knowhere is now Open Source!** We have open-sourced our entire stack for document ingestion, parsing, and agentic RAG. You can now self-host the full platform using [knowhere-self-hosted](https://github.com/Ontos-AI/knowhere-self-hosted). Check out our [Contribution Guide](CONTRIBUTING.md) to get involved! @@ -78,13 +79,15 @@ Knowhere runs in two steps: build memory from documents, then let agents retriev ### Step 2: Agentic Retrieval +Knowhere provides the document-memory substrate; the agent decides how to explore it. +

- Step 2: Agentic Retrieval + Any agent can explore Knowhere Memory through one neutral corpus tool contract

-- **Discover**: Fuse path, content, term, summary, and entity signals for broad first-pass coverage. -- **Navigate**: Use hierarchy-aware MapNav to move from document overviews into the most relevant sections and evidence. -- **Cite Evidence**: Return traceable results with source document, section, source pages, and linked assets. +- **Provide the substrate**: Knowhere exposes one corpus schema and tools for document outlines, structural filters, exact search, fuzzy recall, full reading, assets, and cross-document relationships. +- **Let the agent explore**: Instead of forcing every query through a fixed navigation pipeline, the agent chooses which tools to call, in what order, and how deeply to traverse. +- **Stay agent-neutral**: The same corpus and evidence contract works across built-in agents, MCP clients, models, and orchestration frameworks. Whichever agent explores the memory, Knowhere resolves its references into traceable documents, sections, pages, and linked assets. ## FAQ @@ -141,7 +144,7 @@ Agents using Knowhere outperform those working from raw documents, Markitdown, U - **Vision Page Understanding**: Frontier vision models make complex PDF and PowerPoint content recallable without requiring perfect element-by-element OCR or layout reconstruction. - **Hierarchy-native Memory**: Section nodes preserve document paths, page ranges, summaries, entities, and linked assets instead of returning disconnected chunks. - **Cross-document Memory Graph**: Page-derived typed entities and keywords connect related documents across a namespace. -- **Agentic Retrieval**: MapNav navigates document hierarchies, while classic retrieval combines path, content, and term channels through RRF. +- **Agent-native Retrieval**: Built-in and MCP-connected agents explore the same corpus schema and hierarchy-aware tools; classic retrieval remains available for deterministic top-K search. - **Page-grounded Citations**: Results retain source documents, section paths, page numbers, and rendered visual evidence. ## Supported Formats diff --git a/apps/api/.env.example b/apps/api/.env.example index ef00d0ff..56743634 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -100,8 +100,12 @@ ARK_API_KEY= # Optional retrieval overrides have code defaults. Retrieval is evidence-only: # 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: +# Default path is agent_explore (AGENT_EXPLORE_HARNESS=cursor_sdk). +# Set use_agentic=false for classic map-unit BM25. +# AGENT_EXPLORE_HARNESS=cursor_sdk +# CURSOR_API_KEY= # required when harness is cursor_sdk +# AGENT_EXPLORE_CURSOR_MODEL=composer-2.5 +# Classic BM25 may use Postgres FTS prefilter: # RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT=2000 # File handling defaults diff --git a/apps/api/alembic/env.py b/apps/api/alembic/env.py index 8e3a5d5a..021cc507 100644 --- a/apps/api/alembic/env.py +++ b/apps/api/alembic/env.py @@ -40,21 +40,6 @@ "session", } ) -_AUTOGENERATE_IGNORED_COLUMNS: frozenset[tuple[str, str]] = frozenset( - { - ("document_chunks", "content_search_tsv"), - ("document_chunks", "path_search_tsv"), - } -) - - -def _resolve_table_name(object_: object, compare_to: object | None) -> str | None: - for candidate in (object_, compare_to): - table = getattr(candidate, "table", None) - table_name = getattr(table, "name", None) - if isinstance(table_name, str): - return table_name - return None def include_object( @@ -64,14 +49,10 @@ def include_object( reflected: bool, compare_to: object | None, ) -> bool: - """Exclude externally managed auth tables and generated TSV columns.""" - del reflected + """Exclude externally managed auth tables.""" + del object_, compare_to, reflected if type_ == "table" and isinstance(name, str) and name in _EXTERNALLY_MANAGED_TABLES: return False - if type_ == "column" and isinstance(name, str): - table_name = _resolve_table_name(object_, compare_to) - if table_name is not None and (table_name, name) in _AUTOGENERATE_IGNORED_COLUMNS: - return False return True diff --git a/apps/api/alembic/versions/e4f5a6b7c8d9_drop_chunk_search_tsv_columns.py b/apps/api/alembic/versions/e4f5a6b7c8d9_drop_chunk_search_tsv_columns.py new file mode 100644 index 00000000..94316d6e --- /dev/null +++ b/apps/api/alembic/versions/e4f5a6b7c8d9_drop_chunk_search_tsv_columns.py @@ -0,0 +1,55 @@ +"""Drop unused PostgreSQL FTS columns on document_chunks. + +Retrieval no longer reads ``content_search_tsv`` / ``path_search_tsv``. +They were generated from ``content_search_text`` / ``path_search_text`` and +only served the retired FTS fallback. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op + + +revision: str = "e4f5a6b7c8d9" +down_revision: str | None = "d3e4f5a6b7c8" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +__all__ = [ + "revision", + "down_revision", + "branch_labels", + "depends_on", + "upgrade", + "downgrade", +] + + +def upgrade() -> None: + op.execute("DROP INDEX IF EXISTS idx_chunk_path_search_tsv") + op.execute("DROP INDEX IF EXISTS idx_chunk_content_search_tsv") + op.execute("ALTER TABLE document_chunks DROP COLUMN IF EXISTS path_search_tsv") + op.execute("ALTER TABLE document_chunks DROP COLUMN IF EXISTS content_search_tsv") + + +def downgrade() -> None: + op.execute( + "ALTER TABLE document_chunks ADD COLUMN content_search_tsv TSVECTOR " + "GENERATED ALWAYS AS (to_tsvector('simple', COALESCE(content_search_text, ''))) " + "STORED" + ) + op.execute( + "ALTER TABLE document_chunks ADD COLUMN path_search_tsv TSVECTOR " + "GENERATED ALWAYS AS (to_tsvector('simple', COALESCE(path_search_text, ''))) " + "STORED" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chunk_content_search_tsv " + "ON document_chunks USING GIN (content_search_tsv)" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chunk_path_search_tsv " + "ON document_chunks USING GIN (path_search_tsv)" + ) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index c7d79595..5d84f57f 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -80,8 +80,8 @@ class RetrievalQueryRequest(BaseModel): use_agentic: bool | None = Field( None, description=( - "Map-nav (PLANNER+HARVEST+CONTROL) is the default when unset/true. " - "Set false to force classic 3-channel top-K retrieval." + "Agent explore (cursor_sdk harness by default) when unset/true. " + "Set false to force classic map-unit BM25 top-K retrieval." ), ) conversation_id: str | None = Field( diff --git a/apps/api/app/mcp/dynamic_tools.py b/apps/api/app/mcp/dynamic_tools.py new file mode 100644 index 00000000..efb842f7 --- /dev/null +++ b/apps/api/app/mcp/dynamic_tools.py @@ -0,0 +1,79 @@ +"""Bridge ``shared.services.retrieval.agent_tools.REGISTRY`` onto FastMCP. + +FastMCP's public registration API (``FastMCP.tool`` / ``ToolManager.add_tool``) +only builds a tool's schema by introspecting a Python function's *signature* +(``mcp.server.fastmcp.tools.base.Tool.from_function`` -> ``func_metadata``); +there is no public entry point to register a tool from an already-built JSON +Schema dict, which is what every ``agent_tools.ToolSpec`` carries. Since our +schema is the one already shipped to ``agent_explore`` and meant to be +verbatim-identical across harnesses (see ``CORPUS_SCHEMA.md``), we construct +``Tool`` objects directly instead of round-tripping through a synthetic +Python function signature, and insert them into the tool manager's registry +dict — the same dict ``ToolManager.__init__`` accepts a ``tools=`` list for, +just with no public single-tool equivalent of that constructor path. +""" + +from __future__ import annotations + +from typing import Any, AsyncContextManager, Callable + +from mcp.server.fastmcp import Context, FastMCP +from mcp.server.fastmcp.tools.base import Tool +from mcp.server.fastmcp.utilities.func_metadata import ArgModelBase, FuncMetadata +from pydantic import create_model +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agent_tools import REGISTRY, ToolContext, ToolSpec + +DbFactory = Callable[[], AsyncContextManager[AsyncSession]] + + +def _lenient_arg_model(spec: ToolSpec) -> type[ArgModelBase]: + """Build a permissive pydantic arg model for FastMCP's internal validation. + + The schema actually exposed to MCP clients is ``spec.json_schema`` + (``Tool.parameters``, returned verbatim in ``tools/list`` — see + ``FastMCP.list_tools``); this model only has to satisfy + ``FuncMetadata.call_fn_with_arg_validation`` well enough to forward + whatever the client sent through to ``_dispatch_tool``'s ``**kwargs``. + Every property is optional/``Any`` — each tool already validates its own + required args and reports a caller-facing ``ToolResult.error`` for a + missing one, so duplicating "required" enforcement here would just + produce a less informative MCP-level error instead. + """ + properties = spec.json_schema.get("properties", {}) + fields: dict[str, Any] = {key: (Any, None) for key in properties} + return create_model(f"{spec.name.replace('.', '_')}_Args", __base__=ArgModelBase, **fields) + + +def _make_tool(spec: ToolSpec, *, db_factory: DbFactory) -> Tool: + async def _dispatch_tool( + ctx: Context | None = None, **kwargs: Any + ) -> dict[str, Any]: + from app.mcp.retrieval_server import resolve_mcp_namespace, resolve_mcp_user_id + + namespace = resolve_mcp_namespace(ctx=ctx) + async with db_factory() as db: + user_id = await resolve_mcp_user_id(ctx=ctx, db=db) + tool_ctx = ToolContext(db=db, user_id=user_id, namespace=namespace) + result = await REGISTRY.dispatch(spec.name, tool_ctx, kwargs) + return {"text": result.text, "payload": result.payload, "refs": result.refs, "error": result.error} + + return Tool( + fn=_dispatch_tool, + name=spec.name, + title=None, + description=spec.description, + parameters=spec.json_schema, + fn_metadata=FuncMetadata(arg_model=_lenient_arg_model(spec)), + is_async=True, + context_kwarg="ctx", + annotations=None, + ) + + +def register_corpus_tools(server: FastMCP, *, db_factory: DbFactory) -> None: + """Register every ``agent_tools.REGISTRY`` tool onto ``server``.""" + for spec in REGISTRY.all(): + tool = _make_tool(spec, db_factory=db_factory) + server._tool_manager._tools[tool.name] = tool diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index 81210b86..89845b64 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -10,8 +10,10 @@ from pydantic import Field from sqlalchemy.ext.asyncio import AsyncSession +from app.mcp.dynamic_tools import register_corpus_tools from shared.core.database import get_db_context from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.services.retrieval.agent_tools import load_corpus_schema_text from shared.services.retrieval.app_service import run_retrieval_query from shared.services.retrieval.settings import DEFAULT_TOP_K @@ -85,16 +87,16 @@ def create_retrieval_mcp_server( server = FastMCP( "knowhere-retrieval", instructions=( - "Use this server to search published documents. " - "It returns evidence_text (hierarchical evidence tree), " - "referenced_chunks (structured chunk citations), and " - "decision_trace (navigation decisions). " - "Downstream agents should synthesize answers from evidence_text." + "retrieval.query is a one-shot legacy search tool retained for " + "backward compatibility (see its own tool description below). " + "Prefer the corpus.* tools for exploration; the schema below " + "describes what they operate on.\n\n" + load_corpus_schema_text() ), streamable_http_path=streamable_http_path, stateless_http=True, transport_security=create_public_mcp_transport_security(), ) + register_corpus_tools(server, db_factory=db_factory) @server.tool( name="retrieval.query", diff --git a/apps/api/main.py b/apps/api/main.py index 0af144c8..16912958 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -104,6 +104,18 @@ async def _redis_ping() -> bool: mcp_server = getattr(app.state, "retrieval_mcp_server", None) mcp_session_manager = getattr(mcp_server, "session_manager", None) + from shared.services.retrieval.agent_explore.harness.resolve import ( + resolve_harness_name, + ) + + if resolve_harness_name() == "cursor_sdk" and not os.environ.get( + "CURSOR_API_KEY", "" + ).strip(): + logger.error( + "Default retrieval harness is cursor_sdk but CURSOR_API_KEY is unset; " + "agentic retrieval requests will fail until the key is set" + ) + logger.info("Document API service started!") if mcp_session_manager is not None: async with mcp_session_manager.run(): diff --git a/apps/api/pyproject.toml b/apps/api/pyproject.toml index cb8aaf85..fb67b582 100644 --- a/apps/api/pyproject.toml +++ b/apps/api/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "limits>=5.8,<6", "logfire[celery,fastapi,httpx,sqlalchemy]>=4.25.0", "mcp>=1.27.0", + "cursor-sdk>=1.0.31", ] [dependency-groups] diff --git a/apps/api/scripts/backfill_map_unit_statistics.py b/apps/api/scripts/backfill_map_unit_statistics.py index 86214ef7..1ae6ea79 100644 --- a/apps/api/scripts/backfill_map_unit_statistics.py +++ b/apps/api/scripts/backfill_map_unit_statistics.py @@ -51,7 +51,7 @@ def _bootstrap_python_path() -> None: DocumentMapUnit, DocumentMapUnitIndex, ) -from shared.services.retrieval.nav.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION +from shared.services.retrieval.scoring.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION @dataclass(frozen=True) diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index fe5b732c..d42148d8 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -6,10 +6,11 @@ from uuid import uuid4 import pytest -from httpx import AsyncClient, Response +from httpx import AsyncClient from sqlalchemy import Engine, event, select from shared.models.database.document import DocumentMapUnit +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows from shared.services.retrieval.publication_content import ( replace_document_revision_content, ) @@ -395,100 +396,62 @@ async def test_classic_discovery_returns_empty_for_an_empty_revision_pin( assert result.payload["fused_rows"] == [] -async def test_classic_route_falls_back_for_v1_index_with_excluded_document( +async def test_classic_route_raises_for_v1_index( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], ) -> None: identifier = uuid4().hex[:8] - namespace = f"classic-v1-fallback-{identifier}" - legacy_queries: list[str] = [] - - def capture_legacy_query( - _connection: Any, - _cursor: Any, - statement: str, - _parameters: Any, - _context: Any, - _executemany: bool, - ) -> None: - if "plainto_tsquery('simple'" in statement: - legacy_queries.append(statement) - - event.listen(Engine, "before_cursor_execute", capture_legacy_query) - try: - async with developer_api_client_factory() as api_client: - first = await _publish_document( - namespace=namespace, - source_file_name="legacy-fallback.pdf", - chunks=[ - { - "chunk_id": f"legacy-hit-{identifier}", - "type": "text", - "content": "legacy fallback marker", - "path": "legacy-fallback.pdf/Root/Section/body", - "order": 1, - "metadata": {}, - }, - { - "chunk_id": f"legacy-filler-a-{identifier}", - "type": "text", - "content": "unrelated legacy filler a", - "path": "legacy-fallback.pdf/Root/Section/a", - "order": 2, - "metadata": {}, - }, - { - "chunk_id": f"legacy-filler-b-{identifier}", - "type": "text", - "content": "unrelated legacy filler b", - "path": "legacy-fallback.pdf/Root/Section/b", - "order": 3, - "metadata": {}, - }, - ], - ) - excluded = await _publish_document( - namespace=namespace, - source_file_name="excluded.pdf", - chunks=[ - { - "chunk_id": f"excluded-{identifier}", - "type": "text", - "content": "unrelated filler", - "path": "excluded.pdf/Root/Section/body", - "order": 1, - "metadata": {}, - } - ], - ) - await ContractDatabase.execute( - """ - UPDATE document_map_unit_indexes - SET format_version = 1 - WHERE document_id = :document_id - """, - {"document_id": first["document_id"]}, - ) - response = await api_client.post( + namespace = f"classic-v1-raise-{identifier}" + async with developer_api_client_factory() as api_client: + first = await _publish_document( + namespace=namespace, + source_file_name="legacy-index.pdf", + chunks=[ + { + "chunk_id": f"legacy-hit-{identifier}", + "type": "text", + "content": "legacy index marker", + "path": "legacy-index.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"legacy-filler-a-{identifier}", + "type": "text", + "content": "unrelated filler a", + "path": "legacy-index.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"legacy-filler-b-{identifier}", + "type": "text", + "content": "unrelated filler b", + "path": "legacy-index.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + ], + ) + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET format_version = 1 + WHERE document_id = :document_id + """, + {"document_id": first["document_id"]}, + ) + with pytest.raises(RuntimeError, match="map-unit index is incomplete"): + await api_client.post( "/api/v1/retrieval/query", json={ "namespace": namespace, - "query": "legacy fallback marker", + "query": "legacy index marker", "top_k": 1, "use_agentic": False, - "exclude_document_ids": [excluded["document_id"]], }, ) - finally: - event.remove(Engine, "before_cursor_execute", capture_legacy_query) - - assert response.status_code == 200 - body = cast(dict[str, object], response.json()) - results = cast(list[dict[str, object]], body["results"]) - assert len(results) == 1 - assert results[0]["chunk_id"] == f"legacy-hit-{identifier}" - assert legacy_queries async def test_classic_discovery_preserves_results_before_statistics_backfill( @@ -604,133 +567,95 @@ def result_signature(result: DiscoveryResult) -> list[tuple[Any, ...]]: @pytest.mark.parametrize( "incomplete_index_kind", ["legacy_format", "missing_index", "missing_tokens"] ) -async def test_unfiltered_classic_route_falls_back_when_selective_rows_are_unavailable( +async def test_unfiltered_classic_route_raises_when_index_is_unusable( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], incomplete_index_kind: str, ) -> None: identifier: str = uuid4().hex[:8] - namespace: str = f"classic-token-fallback-{incomplete_index_kind}-{identifier}" - legacy_queries: list[str] = [] - - def capture_legacy_query( - _connection: object, - _cursor: object, - statement: str, - _parameters: object, - _context: object, - _executemany: bool, - ) -> None: - if "plainto_tsquery('simple'" in statement: - legacy_queries.append(statement) - - event.listen(Engine, "before_cursor_execute", capture_legacy_query) - try: - async with developer_api_client_factory() as api_client: - document: dict[str, str] = await _publish_document( - namespace=namespace, - source_file_name="legacy-token-hash.pdf", - chunks=[ - { - "chunk_id": f"legacy-token-hit-{identifier}", - "type": "text", - "content": "legacy token fallback marker", - "path": "legacy-token-hash.pdf/Root/Section/body", - "order": 1, - "metadata": {}, - }, - { - "chunk_id": f"legacy-token-filler-a-{identifier}", - "type": "text", - "content": "unrelated legacy filler a", - "path": "legacy-token-hash.pdf/Root/Section/a", - "order": 2, - "metadata": {}, - }, - { - "chunk_id": f"legacy-token-filler-b-{identifier}", - "type": "text", - "content": "unrelated legacy filler b", - "path": "legacy-token-hash.pdf/Root/Section/b", - "order": 3, - "metadata": {}, - }, - ], + namespace: str = f"classic-unusable-{incomplete_index_kind}-{identifier}" + async with developer_api_client_factory() as api_client: + document: dict[str, str] = await _publish_document( + namespace=namespace, + source_file_name="unusable-index.pdf", + chunks=[ + { + "chunk_id": f"unusable-hit-{identifier}", + "type": "text", + "content": "unusable index marker", + "path": "unusable-index.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"unusable-filler-a-{identifier}", + "type": "text", + "content": "unrelated filler a", + "path": "unusable-index.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"unusable-filler-b-{identifier}", + "type": "text", + "content": "unrelated filler b", + "path": "unusable-index.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + ], + ) + if incomplete_index_kind == "legacy_format": + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET format_version = 1 + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, ) - if incomplete_index_kind == "legacy_format": - await ContractDatabase.execute( - """ - UPDATE document_map_unit_indexes - SET format_version = 1 + elif incomplete_index_kind == "missing_tokens": + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_tokens + WHERE map_unit_id IN ( + SELECT id + FROM document_map_units WHERE document_id = :document_id - """, - {"document_id": document["document_id"]}, - ) - await ContractDatabase.execute( - """ - UPDATE document_map_unit_tokens - SET token_hash = :legacy_token_hash - WHERE map_unit_id IN ( - SELECT id - FROM document_map_units - WHERE document_id = :document_id - ) - """, - { - "document_id": document["document_id"], - "legacy_token_hash": "legacy-token-hash", - }, ) - elif incomplete_index_kind == "missing_tokens": - await ContractDatabase.execute( - """ - DELETE FROM document_map_unit_tokens - WHERE map_unit_id IN ( - SELECT id - FROM document_map_units - WHERE document_id = :document_id - ) - """, - {"document_id": document["document_id"]}, - ) - else: - await ContractDatabase.execute( - """ - DELETE FROM document_map_unit_indexes + """, + {"document_id": document["document_id"]}, + ) + else: + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_indexes + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_tokens + WHERE map_unit_id IN ( + SELECT id + FROM document_map_units WHERE document_id = :document_id - """, - {"document_id": document["document_id"]}, ) - await ContractDatabase.execute( - """ - DELETE FROM document_map_unit_tokens - WHERE map_unit_id IN ( - SELECT id - FROM document_map_units - WHERE document_id = :document_id - ) - """, - {"document_id": document["document_id"]}, - ) - response: Response = await api_client.post( + """, + {"document_id": document["document_id"]}, + ) + with pytest.raises(RuntimeError, match="map-unit index is incomplete"): + await api_client.post( "/api/v1/retrieval/query", json={ "namespace": namespace, - "query": "legacy token fallback marker", + "query": "unusable index marker", "top_k": 1, "use_agentic": False, }, ) - finally: - event.remove(Engine, "before_cursor_execute", capture_legacy_query) - - assert response.status_code == 200 - body = cast(dict[str, object], response.json()) - results = cast(list[dict[str, object]], body["results"]) - assert len(results) == 1 - assert results[0]["chunk_id"] == f"legacy-token-hit-{identifier}" - assert legacy_queries async def test_classic_route_image_filter_scores_only_units_with_images( @@ -856,6 +781,78 @@ async def test_classic_route_image_filter_scores_only_units_with_images( assert results[0]["chunk_type"] == "image" +async def test_connected_hydration_does_not_load_legacy_job_chunks( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"connected-job-{identifier}" + statements: list[str] = [] + + def capture_job_chunk_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if "job_chunks" in statement.lower(): + statements.append(statement) + + async with developer_api_client_factory(): + published = await _publish_document( + namespace=namespace, + source_file_name="connected.pdf", + chunks=[ + { + "chunk_id": "body-connected", + "type": "text", + "content": "body connected evidence", + "path": "connected.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": "asset-connected"}]}, + }, + { + "chunk_id": "asset-connected", + "type": "image", + "content": "asset connected summary", + "path": "images/asset-connected.png", + "order": 2, + "file_path": "images/asset-connected.png", + "metadata": {}, + }, + ], + ) + event.listen(Engine, "before_cursor_execute", capture_job_chunk_query) + try: + async with contract_db_session() as db: + hydrated = await hydrate_connected_target_rows( + db=db, + rows=[ + { + "document_id": published["document_id"], + "job_result_id": published["job_result_id"], + "chunk_id": "body-connected", + "chunk_type": "text", + "chunk_metadata": { + "connect_to": [{"target": "asset-connected"}] + }, + } + ], + exclude_document_ids=[], + exclude_sections=[], + revision_pins={published["document_id"]: published["job_result_id"]}, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_job_chunk_query) + + assert [row["chunk_id"] for row in hydrated] == ["asset-connected"] + assert hydrated[0]["job_id"] == published["job_id"] + assert statements == [] + + def _publish_revision_with_generation_lock( sync_db: Any, *, diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 90dc7e4f..49f676b6 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -245,20 +245,38 @@ async def test_retrieval_should_use_classic_topk_when_agentic_is_false( [], AbstractAsyncContextManager[AsyncClient] ], ) -> None: + from tests.contract.test_retrieval_classic_map_unit_contract import ( + _publish_document, + ) + async with developer_api_client_factory() as api_client: - await _seed_retrieval_document( - user_id="local-dev-user", + await _publish_document( namespace="contract-agentic-only", source_file_name="a.pdf", - section_path="agentic/a", - content="same ranking marker a", - ) - await _seed_retrieval_document( - user_id="local-dev-user", + chunks=[ + { + "chunk_id": "classic-a", + "type": "text", + "content": "same ranking marker a", + "path": "a.pdf/Root/agentic/a", + "order": 1, + "metadata": {}, + } + ], + ) + await _publish_document( namespace="contract-agentic-only", source_file_name="b.pdf", - section_path="agentic/b", - content="same ranking marker b", + chunks=[ + { + "chunk_id": "classic-b", + "type": "text", + "content": "same ranking marker b", + "path": "b.pdf/Root/agentic/b", + "order": 1, + "metadata": {}, + } + ], ) response = await api_client.post( "/api/v1/retrieval/query", @@ -413,62 +431,59 @@ async def test_should_exclude_matching_sections_from_the_response( -def _episode_keeping_chunks( +def _episode_keeping_refs( *, documents: list[dict[str, str]], - evidence_text: str = "mapnav evidence", + notes: str = "", ) -> Any: - """Build a minimal EpisodeResult whose kept_chunks use real seeded chunk_ids.""" - from shared.services.retrieval.nav._compat import AgentStep, Chunk, EpisodeResult + from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult - kept: list[Chunk] = [] - scored: list[tuple[Chunk, float]] = [] + refs: list[dict[str, str]] = [] 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)) + ref = {"document_id": doc["document_id"]} + if doc.get("chunk_id"): + ref["chunk_id"] = doc["chunk_id"] + if doc.get("section_path"): + ref["section_path"] = doc["section_path"] + refs.append(ref) return EpisodeResult( - representation="mapnav", + refs=refs, + notes=notes, 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, - }, + step_index=1, + tool_name="finish", + tool_args={"refs": refs}, + observation_text=notes or "done", + error=None, + elapsed_ms=1, + tokens_used_delta=1, + tokens_used_total=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", + stop_reason="finished", + tokens_used=1, + model_name="test", ) -def _patch_run_nav_episode(monkeypatch: MonkeyPatch, episode: Any) -> None: - def _fake_run_nav_episode(*_args: Any, **_kwargs: Any) -> Any: - return episode +class _FakeHarness: + def __init__(self, episode: Any) -> None: + self._episode = episode + + async def run_episode(self, **_kwargs: Any) -> Any: + return self._episode + +def _patch_harness(monkeypatch: MonkeyPatch, episode: Any) -> None: monkeypatch.setattr( - "shared.services.retrieval.nav.run_nav_episode", - _fake_run_nav_episode, + "shared.services.retrieval.agent_explore.harness.resolve_harness", + lambda: _FakeHarness(episode), ) @pytest.mark.asyncio -async def test_mapnav_retrieval_should_return_seeded_chunk_via_fake_episode( +async def test_agent_explore_retrieval_should_return_seeded_chunk_via_fake_episode( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -477,31 +492,30 @@ async def test_mapnav_retrieval_should_return_seeded_chunk_via_fake_episode( async with developer_api_client_factory() as api_client: target = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-seed", + namespace="contract-explore-seed", source_file_name="target.pdf", section_path="Findings", - content="mapnav seeded EBITDA marker content", + content="explore seeded EBITDA marker content", ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-seed", + namespace="contract-explore-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( + _patch_harness( monkeypatch, - _episode_keeping_chunks( - documents=[target_with_content], - evidence_text="mapnav seeded EBITDA marker content", + _episode_keeping_refs( + documents=[target], + notes="explore seeded EBITDA marker content", ), ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-mapnav-seed", + "namespace": "contract-explore-seed", "query": "EBITDA marker", "top_k": 1, "use_agentic": True, @@ -513,23 +527,17 @@ async def test_mapnav_retrieval_should_return_seeded_chunk_via_fake_episode( 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"] == "mapnav" - assert response_json["stop_reason"] == "completed" + assert response_json["router_used"] == "agent_explore" + assert response_json["stop_reason"] == "finished" 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"} + assert response_json["decision_trace"][-1]["phase"] == "finish" + assert any( + ref.get("chunk_id") == target["chunk_id"] + and ref.get("document_id") == target["document_id"] for ref in referenced_chunks - ] - assert results[0]["content"] == "mapnav seeded EBITDA marker content" + ) + assert results[0]["content"] == "explore seeded EBITDA marker content" assert results[0]["source"] == { "document_id": target["document_id"], "source_file_name": "target.pdf", @@ -538,63 +546,80 @@ async def test_mapnav_retrieval_should_return_seeded_chunk_via_fake_episode( @pytest.mark.asyncio -async def test_mapnav_retrieval_should_not_hydrate_references_outside_request_scope( +async def test_agentic_router_env_mapnav_is_ignored( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], monkeypatch: MonkeyPatch, ) -> None: + monkeypatch.setenv("RETRIEVAL_AGENTIC_ROUTER", "mapnav") async with developer_api_client_factory() as api_client: + target = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-explore-env", + source_file_name="target.pdf", + section_path="Findings", + content="env ignored content", + ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-visible", + namespace="contract-explore-env", + source_file_name="filler.pdf", + section_path="filler/section", + content="unrelated filler content", + ) + _patch_harness( + monkeypatch, + _episode_keeping_refs(documents=[target], notes="env ignored"), + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-explore-env", + "query": "ignored", + "top_k": 1, + "use_agentic": True, + }, + ) + assert response.status_code == 200 + assert cast(dict[str, object], response.json())["router_used"] == "agent_explore" + + +@pytest.mark.asyncio +async def test_agent_explore_should_not_hydrate_references_outside_request_scope( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + async with developer_api_client_factory() as api_client: + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-explore-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-mapnav-visible", + namespace="contract-explore-visible", source_file_name="visible-filler.pdf", section_path="visible/filler", content="visible scoped filler content", ) foreign = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-foreign", + namespace="contract-explore-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.nav_bridge.build_referenced_chunks", - _fake_bridge, - ) + _patch_harness(monkeypatch, _episode_keeping_refs(documents=[foreign])) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-mapnav-visible", + "namespace": "contract-explore-visible", "query": "visible", "top_k": 1, "use_agentic": True, @@ -603,13 +628,13 @@ def _fake_bridge(_episode: Any, _snapshot: Any) -> tuple[list[dict[str, Any]], d assert response.status_code == 200 response_json = cast(dict[str, object], response.json()) - assert response_json["router_used"] == "mapnav" + assert response_json["router_used"] == "agent_explore" assert response_json["referenced_chunks"] == [] assert response_json["results"] == [] @pytest.mark.asyncio -async def test_mapnav_retrieval_should_drop_references_with_mismatched_section_path( +async def test_agent_explore_should_drop_references_with_mismatched_section_path( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -618,47 +643,28 @@ async def test_mapnav_retrieval_should_drop_references_with_mismatched_section_p async with developer_api_client_factory() as api_client: visible = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-section-mismatch", + namespace="contract-explore-section-mismatch", source_file_name="visible.pdf", section_path="visible/section", content="visible scoped content", ) await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-section-mismatch", + namespace="contract-explore-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.nav_bridge.build_referenced_chunks", - _fake_bridge, - ) + mismatched = { + "document_id": visible["document_id"], + "section_path": "wrong/section/path", + } + _patch_harness(monkeypatch, _episode_keeping_refs(documents=[mismatched])) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-mapnav-section-mismatch", + "namespace": "contract-explore-section-mismatch", "query": "visible", "top_k": 1, "use_agentic": True, @@ -667,13 +673,13 @@ def _fake_bridge(_episode: Any, _snapshot: Any) -> tuple[list[dict[str, Any]], d assert response.status_code == 200 response_json = cast(dict[str, object], response.json()) - assert response_json["router_used"] == "mapnav" + assert response_json["router_used"] == "agent_explore" assert response_json["referenced_chunks"] == [] assert response_json["results"] == [] @pytest.mark.asyncio -async def test_mapnav_retrieval_should_fail_when_final_hydration_db_fails( +async def test_agent_explore_should_fail_when_final_hydration_db_fails( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -685,26 +691,21 @@ async def fail_final_hydration(**_kwargs: object) -> object: async with developer_api_client_factory() as api_client: visible = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-hydration-failure", + namespace="contract-explore-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-mapnav-hydration-failure", + namespace="contract-explore-hydration-failure", source_file_name="filler.pdf", section_path="filler/section", content="filler content", ) from shared.services.retrieval.execution import routes as retrieval_routes - _patch_run_nav_episode( - monkeypatch, - _episode_keeping_chunks( - documents=[{**visible, "content": "visible scoped content"}] - ), - ) + _patch_harness(monkeypatch, _episode_keeping_refs(documents=[visible])) monkeypatch.setattr( retrieval_routes, "resolve_workflow_references", @@ -718,7 +719,7 @@ async def fail_final_hydration(**_kwargs: object) -> object: await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-mapnav-hydration-failure", + "namespace": "contract-explore-hydration-failure", "query": "visible", "top_k": 1, "use_agentic": True, @@ -727,7 +728,7 @@ async def fail_final_hydration(**_kwargs: object) -> object: @pytest.mark.asyncio -async def test_mapnav_should_preserve_same_chunk_id_across_documents( +async def test_agent_explore_should_preserve_same_chunk_id_across_documents( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -738,7 +739,7 @@ async def test_mapnav_should_preserve_same_chunk_id_across_documents( async with developer_api_client_factory() as api_client: first = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-shared-chunk", + namespace="contract-explore-shared-chunk", source_file_name="first.pdf", section_path="shared/first", content="first shared reference content", @@ -746,34 +747,29 @@ async def test_mapnav_should_preserve_same_chunk_id_across_documents( ) second_doc = await _seed_retrieval_document( user_id="local-dev-user", - namespace="contract-mapnav-shared-chunk", + namespace="contract-explore-shared-chunk", source_file_name="second.pdf", section_path="shared/second-host", content="host content for second document", ) second = await _seed_retrieval_chunk_for_existing_document( user_id="local-dev-user", - namespace="contract-mapnav-shared-chunk", + namespace="contract-explore-shared-chunk", document=second_doc, section_path="shared/second", content="second shared reference content", chunk_id=shared_chunk_id, ) - _patch_run_nav_episode( + _patch_harness( monkeypatch, - _episode_keeping_chunks( - documents=[ - {**first, "content": "first shared reference content"}, - {**second, "content": "second shared reference content"}, - ] - ), + _episode_keeping_refs(documents=[first, second]), ) response = await api_client.post( "/api/v1/retrieval/query", json={ - "namespace": "contract-mapnav-shared-chunk", + "namespace": "contract-explore-shared-chunk", "query": "shared reference", "top_k": 1, "use_agentic": True, @@ -785,7 +781,7 @@ async def test_mapnav_should_preserve_same_chunk_id_across_documents( 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"] == "mapnav" + assert response_json["router_used"] == "agent_explore" assert len(referenced_chunks) == 2 assert {ref["document_id"] for ref in referenced_chunks} == { first["document_id"], diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 46a347c8..3b73cf4f 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -95,8 +95,12 @@ ARK_API_KEY= # Optional retrieval overrides have code defaults. Retrieval is evidence-only: # 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: +# Default path is agent_explore (AGENT_EXPLORE_HARNESS=cursor_sdk). +# Set use_agentic=false for classic map-unit BM25. +# AGENT_EXPLORE_HARNESS=cursor_sdk +# CURSOR_API_KEY= # required when harness is cursor_sdk +# AGENT_EXPLORE_CURSOR_MODEL=composer-2.5 +# Classic BM25 may use Postgres FTS prefilter: # RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT=2000 # Required for specific features: billing and analytics diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 2b99effb..a7bd5247 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -100,6 +100,14 @@ def _ignore_gevent_cancellation( @worker_shutdown.connect def shutdown_worker(**kwargs: object) -> None: """Clean up shared resources on worker shutdown.""" + try: + from app.services.document_agent.tools.ocr_pages import shutdown_ocr_runner + + shutdown_ocr_runner() + logger.info("Worker OCR runner stopped") + except Exception as exc: + logger.warning(f"Worker OCR runner cleanup failed: {exc}") + try: stop_worker_heartbeat() logger.info("Worker heartbeat stopped") diff --git a/apps/worker/app/services/document_agent/tools/ocr_pages.py b/apps/worker/app/services/document_agent/tools/ocr_pages.py index 80565118..b1a29db4 100644 --- a/apps/worker/app/services/document_agent/tools/ocr_pages.py +++ b/apps/worker/app/services/document_agent/tools/ocr_pages.py @@ -2,17 +2,404 @@ from __future__ import annotations +import atexit +import multiprocessing import time -from typing import Any +from multiprocessing.connection import Connection +from multiprocessing.process import BaseProcess +from threading import RLock +from typing import TYPE_CHECKING, Any, NotRequired, TypedDict, cast + +from loguru import logger from app.services.document_agent.manifest import ToolContext, ToolResult from app.services.document_agent.pdf_text import PageTextBands -from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( - run_in_child_process, - worker, -) from app.services.document_agent.registry import has_page_features, register_tool from app.services.document_agent.visual import render_pages +from shared.core.exceptions.domain_exceptions import PDFParsingException, TimeoutException + +if TYPE_CHECKING: + from gevent.lock import Semaphore + from gevent.threadpool import ThreadPool + + +OCR_TIMEOUT_SECONDS = 1800 +OCR_INTRA_OP_THREADS = 3 +OCR_INTER_OP_THREADS = 1 +OCR_QUEUE_POLL_INTERVAL_SECONDS = 0.1 +OCR_CHILD_JOIN_TIMEOUT_SECONDS = 5 +OCR_RUNNER_CONTEXT = multiprocessing.get_context("spawn") + + +class _OcrRequest(TypedDict): + page: int + image_path: str + + +class _OcrResponse(TypedDict): + ok: bool + page: NotRequired[int] + lines: NotRequired[list[dict[str, Any]]] + error_type: NotRequired[str] + error_msg: NotRequired[str] + + +def _run_ocr_process( + request_connection: Connection, + response_connection: Connection, +) -> None: + """Serve sequential OCR requests while keeping one model instance alive.""" + try: + from rapidocr_onnxruntime import RapidOCR + + engine = RapidOCR( + intra_op_num_threads=OCR_INTRA_OP_THREADS, + inter_op_num_threads=OCR_INTER_OP_THREADS, + ) + except Exception as exc: + response_connection.send( + { + "ok": False, + "error_type": type(exc).__name__, + "error_msg": str(exc), + } + ) + _close_child_connections(request_connection, response_connection) + return + + try: + while True: + request = request_connection.recv() + if request is None: + return + page = int(request["page"]) + image_path = str(request["image_path"]) + try: + result, _elapse = engine(image_path) + lines: list[dict[str, Any]] = [] + for item in result or []: + lines.append( + { + "box": _line_box(item), + "text": _line_text(item), + "score": _line_score(item), + } + ) + response_connection.send({"ok": True, "page": page, "lines": lines}) + except Exception as exc: + response_connection.send( + { + "ok": False, + "page": page, + "error_type": type(exc).__name__, + "error_msg": str(exc), + } + ) + finally: + _close_child_connections(request_connection, response_connection) + + +def _close_child_connections( + request_connection: Connection, + response_connection: Connection, +) -> None: + request_connection.close() + response_connection.close() + + +def _close_parent_connection(connection: Connection) -> None: + try: + connection.close() + except Exception as exc: + logger.debug("Ignoring error while closing parent connection: {}", exc) + + +class _OcrRunner: + """Own one persistent OCR child and serialize page requests through it.""" + + def __init__(self) -> None: + self._process: BaseProcess | None = None + self._request_connection: Connection | None = None + self._response_connection: Connection | None = None + self._executor: ThreadPool | None = None + self._slot: Semaphore | None = None + self._lock = RLock() + + def run_page(self, page: int, image_path: str, *, timeout: int) -> _OcrResponse: + deadline = time.monotonic() + timeout + slot = self._get_slot() + remaining = deadline - time.monotonic() + if remaining <= 0 or not slot.acquire(timeout=remaining): + raise TimeoutException( + internal_message=( + f"OCR request waited more than {timeout}s for the shared runner: " + f"page={page}" + ), + retry_after=30, + ) + + task = None + try: + executor = self._get_executor() + task = executor.spawn( + self._run_page_blocking, + page, + image_path, + deadline, + ) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutException( + internal_message=( + f"OCR request timed out before execution: page={page}" + ), + retry_after=30, + ) + raw_result = task.get(timeout=remaining) + if not isinstance(raw_result, dict): + raise PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_FAILED", + internal_message=f"OCR runner returned an invalid result: page={page}", + ) + return cast(_OcrResponse, raw_result) + except BaseException as exc: + self._cancel_active_process() + if task is not None: + try: + task.get(timeout=OCR_CHILD_JOIN_TIMEOUT_SECONDS) + except BaseException as join_exc: + # Best-effort cleanup: preserve the original exception `exc`. + logger.debug( + "Ignoring OCR task join failure during cancellation " + "(page={}, join_timeout={}s): {}: {}", + page, + OCR_CHILD_JOIN_TIMEOUT_SECONDS, + type(join_exc).__name__, + join_exc, + ) + from gevent.timeout import Timeout + + if isinstance(exc, Timeout): + raise TimeoutException( + internal_message=( + f"OCR request timed out after {timeout}s: page={page}" + ), + retry_after=30, + ) from exc + raise + finally: + slot.release() + + def close(self) -> None: + process, request_connection, response_connection = self._detach_process() + self._stop_process(process, request_connection, response_connection) + + with self._lock: + executor = self._executor + self._executor = None + if executor is not None: + executor.kill() + executor.join() + + def _get_executor(self) -> ThreadPool: + if self._executor is None: + from gevent.threadpool import ThreadPool + + self._executor = ThreadPool(maxsize=1) + return self._executor + + def _get_slot(self) -> Semaphore: + if self._slot is None: + from gevent.lock import Semaphore + + self._slot = Semaphore(1) + return self._slot + + def _run_page_blocking( + self, + page: int, + image_path: str, + deadline: float, + ) -> _OcrResponse: + self._ensure_process() + request_connection = self._request_connection + response_connection = self._response_connection + process = self._process + if request_connection is None or response_connection is None or process is None: + raise RuntimeError("OCR runner failed to initialize") + + request: _OcrRequest = {"page": page, "image_path": image_path} + request_connection.send(request) + try: + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + self._stop_active_process( + process, request_connection, response_connection + ) + raise TimeoutException( + internal_message=( + f"OCR child timed out after {OCR_TIMEOUT_SECONDS}s: " + f"page={page}, pid={process.pid}" + ), + retry_after=30, + ) + try: + if response_connection.poll( + timeout=min(remaining, OCR_QUEUE_POLL_INTERVAL_SECONDS) + ): + result = response_connection.recv() + break + if process.is_alive(): + continue + self._stop_active_process( + process, request_connection, response_connection + ) + raise PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_CRASH", + internal_message=( + f"OCR child exited without a result: page={page}, " + f"pid={process.pid}" + ), + ) + except Exception as exc: + self._stop_active_process( + process, request_connection, response_connection + ) + raise PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_FAILED", + internal_message=( + f"OCR child response queue failed on page {page}: " + f"{type(exc).__name__}: {exc}" + ), + ) from exc + except BaseException: + raise + + if not isinstance(result, dict): + self._stop_active_process( + process, request_connection, response_connection + ) + raise PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_FAILED", + internal_message=f"OCR child returned an invalid response: page={page}", + ) + if not result.get("ok"): + error_type = result.get("error_type", "UnknownError") + error_msg = result.get("error_msg", "unknown OCR child failure") + raise PDFParsingException( + user_message="Failed to process your document. Please try again.", + reason="SUBPROCESS_FAILED", + internal_message=( + f"OCR child failed on page {page}: {error_type}: {error_msg}" + ), + ) + return cast(_OcrResponse, result) + + def _ensure_process(self) -> None: + with self._lock: + if self._process is not None and self._process.is_alive(): + return + process = self._process + request_connection = self._request_connection + response_connection = self._response_connection + self._process = None + self._request_connection = None + self._response_connection = None + self._stop_process(process, request_connection, response_connection) + + child_request_connection, request_connection = OCR_RUNNER_CONTEXT.Pipe( + duplex=False + ) + response_connection, child_response_connection = OCR_RUNNER_CONTEXT.Pipe( + duplex=False + ) + process = OCR_RUNNER_CONTEXT.Process( + target=_run_ocr_process, + args=(child_request_connection, child_response_connection), + ) + self._request_connection = request_connection + self._response_connection = response_connection + self._process = process + try: + process.start() + except BaseException: + self._process = None + self._request_connection = None + self._response_connection = None + self._stop_process( + process, + request_connection, + response_connection, + child_request_connection, + child_response_connection, + ) + raise + child_request_connection.close() + child_response_connection.close() + logger.info("[ocr-runner] started persistent OCR child pid={}", process.pid) + + def _detach_process( + self, + ) -> tuple[ + BaseProcess | None, + Connection | None, + Connection | None, + ]: + with self._lock: + process = self._process + request_connection = self._request_connection + response_connection = self._response_connection + self._process = None + self._request_connection = None + self._response_connection = None + return process, request_connection, response_connection + + def _cancel_active_process(self) -> None: + process, request_connection, response_connection = self._detach_process() + self._stop_process(process, request_connection, response_connection) + + def _stop_active_process( + self, + process: BaseProcess, + request_connection: Connection, + response_connection: Connection, + ) -> None: + with self._lock: + if self._process is not process: + return + self._process = None + self._request_connection = None + self._response_connection = None + self._stop_process(process, request_connection, response_connection) + + @staticmethod + def _stop_process( + process: BaseProcess | None, + request_connection: Connection | None, + response_connection: Connection | None, + *child_connections: Connection, + ) -> None: + if process is not None and process.pid is not None: + if process.is_alive(): + process.kill() + process.join(timeout=OCR_CHILD_JOIN_TIMEOUT_SECONDS) + for connection in (request_connection, response_connection, *child_connections): + if connection is not None: + _close_parent_connection(connection) + + +_OCR_RUNNER = _OcrRunner() +atexit.register(_OCR_RUNNER.close) + + +def shutdown_ocr_runner() -> None: + """Stop the worker-local OCR runner during Celery shutdown.""" + _OCR_RUNNER.close() def _line_text(item: Any) -> str: @@ -36,28 +423,6 @@ def _line_score(item: Any) -> float: return 0.0 -@worker -def _run_ocr_worker(queue: Any, page_paths: dict[int, str]) -> None: - """Run the local OCR model outside the heartbeat-bearing worker process.""" - from rapidocr_onnxruntime import RapidOCR - - engine = RapidOCR(intra_op_num_threads=2, inter_op_num_threads=1) - page_lines: dict[int, list[dict[str, Any]]] = {} - for page, image_path in page_paths.items(): - lines: list[dict[str, Any]] = [] - result, _elapse = engine(image_path) - for item in result or []: - lines.append( - { - "box": _line_box(item), - "text": _line_text(item), - "score": _line_score(item), - } - ) - page_lines[page] = lines - queue.put({"ok": True, "page_lines": page_lines}) - - @register_tool( name="ocr.pages", description="Run RapidOCR on specified pages and return positioned text lines.", @@ -96,14 +461,19 @@ def ocr_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: if item.get("page") is not None and item.get("png_path") } - page_paths = { - page: image_path for page, image_path in png_by_page.items() if page in pages - } - result = run_in_child_process(_run_ocr_worker, page_paths, timeout=300) - page_lines = { - int(page): list(lines) - for page, lines in (result.get("page_lines") or {}).items() - } + page_lines: dict[int, list[dict[str, Any]]] = {} + for page in pages: + image_path = png_by_page.get(page) + if not image_path: + page_lines[page] = [] + continue + result = _OCR_RUNNER.run_page( + page, + image_path, + timeout=OCR_TIMEOUT_SECONDS, + ) + result_page = int(result.get("page") or page) + page_lines[result_page] = list(result.get("lines") or []) page_texts: dict[int, str] = {} page_bands: dict[int, PageTextBands] = {} for page in pages: diff --git a/apps/worker/pyproject.toml b/apps/worker/pyproject.toml index 021eb6f9..8cb35cdb 100644 --- a/apps/worker/pyproject.toml +++ b/apps/worker/pyproject.toml @@ -30,6 +30,13 @@ dependencies = [ "rapidocr-onnxruntime>=1.4.4", ] +# Production retrieval runs in apps/api, where cursor-sdk is a base +# dependency. This extra is only for worker debug/eval scripts. +[project.optional-dependencies] +cursor-harness = [ + "cursor-sdk>=1.0.31", +] + [dependency-groups] dev = [ "fakeredis[lua]>=2.31.0", diff --git a/apps/worker/scripts/debug_agent_explore_episode.py b/apps/worker/scripts/debug_agent_explore_episode.py new file mode 100644 index 00000000..c85c531f --- /dev/null +++ b/apps/worker/scripts/debug_agent_explore_episode.py @@ -0,0 +1,92 @@ +"""Audit-only: dump every step of one ``agent_explore`` episode. + +Prints, for each LLM turn / tool call: tool name, args, elapsed ms, +tokens_used_delta (turn-level, attributed to the first tool step — see +``harness/openai_harness.py``), tokens_used_total (cumulative), observation +length (chars sent back into the LLM's context), and error. Also prints the +final ``EpisodeResult`` (refs/notes/stop_reason). + +Read-only diagnostic; does not modify any behavior. + +Usage: + cd apps/worker + uv run python scripts/debug_agent_explore_episode.py --query-id q04 + uv run python scripts/debug_agent_explore_episode.py --query "..." --token-limit 200000 + uv run python scripts/debug_agent_explore_episode.py --query-id q04 --harness cursor_sdk +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +from pathlib import Path + +from dotenv import load_dotenv + +_SCRIPT_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SCRIPT_DIR.parents[1] +sys.path.insert(0, str(_REPO_ROOT / "packages" / "shared-python")) +sys.path.insert(0, str(_SCRIPT_DIR.parent)) + +load_dotenv(_SCRIPT_DIR.parent / ".env") +os.environ.setdefault("LOCAL_DEBUG", "0") +os.environ.setdefault("LLM_MOCK_ENABLED", "false") + +FIXTURE_PATH = _SCRIPT_DIR / "fixtures" / "changheba_archive_eval_queries.json" + + +async def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--query-id", default=None) + parser.add_argument("--query", default=None) + parser.add_argument("--user-id", default="debug_local_user") + parser.add_argument("--namespace", default="default") + parser.add_argument("--token-limit", type=int, default=None) + parser.add_argument("--harness", default=None, choices=["openai", "cursor_sdk"]) + args = parser.parse_args() + + query = args.query + if args.query_id: + fixture = json.loads(FIXTURE_PATH.read_text(encoding="utf-8")) + match = next(q for q in fixture["queries"] if q["id"] == args.query_id) + query = match["query"] + if not query: + raise SystemExit("need --query or --query-id") + + from shared.core.database import get_db_context + from shared.services.retrieval.agent_explore.budget import EpisodeBudget + from shared.services.retrieval.agent_explore.harness import resolve_harness + + budget = EpisodeBudget(token_limit=args.token_limit) if args.token_limit else EpisodeBudget() + harness = resolve_harness(args.harness) + + print(f"query: {query!r}") + episode = await harness.run_episode( + db_factory=get_db_context, + user_id=args.user_id, + namespace=args.namespace, + query=query, + budget=budget, + ) + + print(f"\nstop_reason={episode.stop_reason} tokens_used={episode.tokens_used} " + f"model={episode.model_name}") + print(f"final refs ({len(episode.refs)}): {json.dumps(episode.refs, ensure_ascii=False)}") + print(f"final notes: {episode.notes!r}") + print(f"\n{'#':>3} {'tool':<28} {'ms':>6} {'delta':>7} {'total':>7} {'obs_chars':>9} err") + for step in episode.steps: + args_preview = json.dumps(step.tool_args, ensure_ascii=False)[:80] + err = step.error or "" + print( + f"{step.step_index:>3} {step.tool_name or '(none)':<28} " + f"{step.elapsed_ms:>6} {step.tokens_used_delta:>7} " + f"{step.tokens_used_total:>7} {len(step.observation_text):>9} {err}" + ) + print(f" args: {args_preview}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/apps/worker/scripts/debug_finish_raw_args.py b/apps/worker/scripts/debug_finish_raw_args.py new file mode 100644 index 00000000..1d6acfa6 --- /dev/null +++ b/apps/worker/scripts/debug_finish_raw_args.py @@ -0,0 +1,66 @@ +"""Audit-only: monkeypatch to see the RAW ``finish`` tool_call.function.arguments +string the LLM actually sent, before ``_safe_json_loads`` parses it. Does not +modify any file — patches the imported module object in this process only. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from pathlib import Path + +from dotenv import load_dotenv + +_SCRIPT_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SCRIPT_DIR.parents[1] +sys.path.insert(0, str(_REPO_ROOT / "packages" / "shared-python")) +sys.path.insert(0, str(_SCRIPT_DIR.parent)) + +load_dotenv(_SCRIPT_DIR.parent / ".env") +os.environ.setdefault("LOCAL_DEBUG", "0") +os.environ.setdefault("LLM_MOCK_ENABLED", "false") + +FIXTURE_PATH = _SCRIPT_DIR / "fixtures" / "changheba_archive_eval_queries.json" + + +async def main() -> None: + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--query-id", required=True) + args = parser.parse_args() + + fixture = json.loads(FIXTURE_PATH.read_text(encoding="utf-8")) + match = next(q for q in fixture["queries"] if q["id"] == args.query_id) + query = match["query"] + + from shared.services.retrieval.agent_explore.harness import openai_harness + + original_safe_json_loads = openai_harness._safe_json_loads + + def _spy(raw): + print(f"[SPY] raw arguments received: {raw!r}") + return original_safe_json_loads(raw) + + openai_harness._safe_json_loads = _spy + + from shared.core.database import get_db_context + from shared.services.retrieval.agent_explore.budget import EpisodeBudget + + print(f"query: {query!r}") + result = await openai_harness.OpenAIHarness().run_episode( + db_factory=get_db_context, + user_id="debug_local_user", + namespace="default", + query=query, + budget=EpisodeBudget(), + ) + print(f"\nfinal refs: {result.refs}") + print(f"final notes: {result.notes!r}") + print(f"stop_reason: {result.stop_reason}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/apps/worker/scripts/fixtures/changheba_archive_eval_queries.json b/apps/worker/scripts/fixtures/changheba_archive_eval_queries.json new file mode 100644 index 00000000..b1c640b8 --- /dev/null +++ b/apps/worker/scripts/fixtures/changheba_archive_eval_queries.json @@ -0,0 +1,93 @@ +{ + "version": "1.0", + "source_docx": "/Users/wuchengke/Desktop/temp/test_docs/yideng/zh_档案知识库测试样例.docx", + "corpus": { + "user_id": "debug_local_user", + "namespace": "default", + "documents": [ + "zh_四川省大渡河长河坝水电站可行性研究报告1.pdf", + "zh_四川省大渡河长河坝水电站可行性研究报告2.pdf" + ] + }, + "queries": [ + { + "id": "q01", + "label": "roles_wangrenkun", + "query": "王仁坤作为不同角色参与了哪些设计成果的编制？", + "answer_notes": [ + "作为总工程师：报告1综合说明封面、报告2水文泥沙封面", + "作为审查人：报告1审查意见汇编（报告2水文泥沙正文未检出审查人署名）" + ], + "expected_keywords": ["王仁坤", "总工程师", "审查"] + }, + { + "id": "q02", + "label": "hydro_stations", + "query": "长河坝水电站上下游分别设置了几个水文站？水文站的水位观测、流量测验情况分别是怎样的？", + "answer_notes": [ + "报告2水文泥沙2.3：干流3站（上大金/丹巴、下泸定），支流金汤、康定", + "水位枯2/汛4（涨水加测），流量流速仪与浮标结合" + ], + "expected_keywords": ["水文", "丹巴", "泸定", "水位", "流量"] + }, + { + "id": "q03", + "label": "project_status_power_scope", + "query": "长河坝水电站的工程地位和作用是什么？供电范围是什么？", + "answer_notes": [ + "报告1综合说明1.4：装机约2600MW，大渡河重点开发工程，川电东送/西电东送", + "供电范围：四川主网，华中、华东电网" + ], + "expected_keywords": ["2600", "川电东送", "华东", "华中"] + }, + { + "id": "q04", + "label": "survey_design_challenges", + "query": "长河坝水电站勘察设计过程中遇到的难点有哪些？采取了哪些关键技术解决难点？", + "answer_notes": [ + "难点：高地震烈度、深厚覆盖层、世界级砾石土心墙堆石坝", + "关键技术：砾石土直心墙、混凝土防渗墙、覆盖层处理、抗震措施" + ], + "expected_keywords": ["砾石土", "防渗墙", "地震", "覆盖层"] + }, + { + "id": "q05", + "label": "installed_capacity", + "query": "长河坝水电站在可行性研究阶段推荐的水电站装机容量是多少？安装几台机组？单机容量是多少？", + "answer_notes": [ + "报告1综合说明1.4.11/1.4.13：推荐260万kW，4台，单机65万kW混流式" + ], + "expected_keywords": ["260", "4", "65", "万 kW"] + }, + { + "id": "q06", + "label": "reservoir_operation", + "query": "水库和电站运行方式是怎样的？", + "answer_notes": [ + "报告1综合说明1.4.14：日/周调节，1690m-1680m，特枯可降至1650m", + "黄金坪投产后自由调峰；投产前须保证5%生态流量" + ], + "expected_keywords": ["1690", "1680", "调峰", "1650"] + }, + { + "id": "q07", + "label": "project_class_standards", + "query": "长河坝水电站的工程等别和设计标准是怎样的？", + "answer_notes": [ + "报告1综合说明1.6.1：一等大（1）型，1级建筑物", + "洪水：挡泄1000年一遇7650，厂房200年6670；地震：壅水Ⅸ度359gal，非壅水Ⅷ度222gal" + ], + "expected_keywords": ["一等大", "1000", "7650", "Ⅸ"] + }, + { + "id": "q08", + "label": "dam_type_selection", + "query": "长河坝水电站有几种坝型选择？最终选择的是哪种？依据是什么？", + "answer_notes": [ + "报告1综合说明1.6.2：比选3种（直心墙、斜心墙、沥青混凝土心墙）", + "最终选择砾石土直心墙堆石坝；沥青心墙接头复杂、经验少" + ], + "expected_keywords": ["砾石土直心墙", "沥青", "三种", "斜心墙"] + } + ] +} diff --git a/apps/worker/scripts/run_agentic_router_eval.py b/apps/worker/scripts/run_agentic_router_eval.py new file mode 100644 index 00000000..328503b4 --- /dev/null +++ b/apps/worker/scripts/run_agentic_router_eval.py @@ -0,0 +1,305 @@ +"""Smoke-run agent_explore on a fixed query set. + +Reads ``fixtures/changheba_archive_eval_queries.json`` (sourced from +``zh_档案知识库测试样例.docx``) and runs each query through +``run_retrieval_route`` directly, bypassing the Redis result cache. + +Usage: + cd apps/worker + uv run python scripts/run_agentic_router_eval.py + uv run python scripts/run_agentic_router_eval.py --query-id q05 + uv run python scripts/run_agentic_router_eval.py --use-agentic false +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any + +from dotenv import load_dotenv + +_SCRIPT_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SCRIPT_DIR.parents[1] +_SHARED_PYTHON = _REPO_ROOT / "packages" / "shared-python" +sys.path.insert(0, str(_SHARED_PYTHON)) +sys.path.insert(0, str(_SCRIPT_DIR.parent)) + +load_dotenv(_SCRIPT_DIR.parent / ".env", override=True) +os.environ.setdefault("LOCAL_DEBUG", "0") +os.environ.setdefault("LLM_MOCK_ENABLED", "false") + +FIXTURE_PATH = _SCRIPT_DIR / "fixtures" / "changheba_archive_eval_queries.json" +_MODES = ("classic", "agent_explore") +TOP_K = 10 + + +def _count_llm_steps(decision_trace: list[dict[str, Any]], mode: str) -> int: + del mode + return sum( + 1 + for step in decision_trace + if step.get("phase") in ("tool_call", "finish") + or (step.get("decision") or {}).get("action") not in ("", "no_tool_call", None) + ) + + +def _keyword_hits(evidence_text: str, keywords: list[str]) -> tuple[int, list[str]]: + text = evidence_text or "" + hits = [kw for kw in keywords if kw and kw in text] + return len(hits), hits + + +@dataclass +class RunMetrics: + query_id: str + router: str # "classic" | "agent_explore" + total_ms: int + router_used: str + stop_reason: str + refs: int + results: int + evidence_chars: int + llm_steps: int + keyword_hits: int + keyword_total: int + matched_keywords: list[str] + error: str | None = None + + +async def _run_one( + *, + user_id: str, + namespace: str, + query: str, + query_id: str, + router: str, + use_agentic: bool, + expected_keywords: list[str], +) -> RunMetrics: + from dataclasses import replace + + from shared.core.database import get_db_context + from shared.services.retrieval.execution.plan import project_public_retrieval_response + from shared.services.retrieval.execution.query_request import RetrievalQuery + from shared.services.retrieval.execution.revision_pins import ( + capture_revision_pins, + is_revision_generation_stable, + ) + from shared.services.retrieval.execution.routes import run_retrieval_route + from shared.services.retrieval.settings import INTERNAL_RECALL_K_MULTIPLIER + + started = time.monotonic() + error: str | None = None + response: dict[str, Any] = {} + try: + async with get_db_context() as db: + request = RetrievalQuery.from_parameters( + db=db, + user_id=user_id, + namespace=namespace, + query=query, + top_k=TOP_K, + exclude_document_ids=[], + exclude_sections=[], + use_agentic=use_agentic, + ) + revision_pins = await capture_revision_pins( + db, user_id=user_id, namespace=namespace + ) + if not await is_revision_generation_stable( + db, + user_id=user_id, + namespace=namespace, + pins=revision_pins, + ): + revision_pins = await capture_revision_pins( + db, user_id=user_id, namespace=namespace + ) + effective_recall_k = ( + request.internal_recall_k + if request.internal_recall_k is not None + else TOP_K * INTERNAL_RECALL_K_MULTIPLIER + ) + context = replace( + request.build_route_context(), + revision_pins=revision_pins, + effective_recall_k=effective_recall_k, + ) + outcome = await run_retrieval_route(context) + response = await project_public_retrieval_response(outcome.response) + except Exception as exc: # noqa: BLE001 - eval runner must continue + error = f"{type(exc).__name__}: {exc}" + + elapsed_ms = int((time.monotonic() - started) * 1000) + evidence = str(response.get("evidence_text") or "") + decision_trace = response.get("decision_trace") or [] + hits, matched = _keyword_hits(evidence, expected_keywords) + return RunMetrics( + query_id=query_id, + router=router, + total_ms=elapsed_ms, + router_used=str(response.get("router_used") or router), + stop_reason=str(response.get("stop_reason") or ""), + refs=len(response.get("referenced_chunks") or []), + results=len(response.get("results") or []), + evidence_chars=len(evidence), + llm_steps=_count_llm_steps(decision_trace, router), + keyword_hits=hits, + keyword_total=len(expected_keywords), + matched_keywords=matched, + error=error, + ) + + +def _render_markdown( + fixture: dict[str, Any], + runs: list[RunMetrics], + output_dir: Path, +) -> str: + by_key = {(r.query_id, r.router): r for r in runs} + lines = [ + "# Agentic Router Eval (Phase 4)\n", + f"Fixture: `{FIXTURE_PATH.name}`\n", + f"Corpus: `{fixture['corpus']['namespace']}` / " + f"{len(fixture['corpus']['documents'])} documents\n", + f"Generated: {datetime.now().isoformat(timespec='seconds')}\n", + f"Output dir: `{output_dir}`\n", + "\n## Summary\n", + "| Q | Query (short) | Router | ms | LLM steps | refs | kw hit | stop |\n", + "|---|---|---|---:|---:|---:|---:|---|\n", + ] + for item in fixture["queries"]: + qid = item["id"] + short = item["query"][:28] + ("…" if len(item["query"]) > 28 else "") + for router in _MODES: + r = by_key.get((qid, router)) + if r is None: + continue + kw = f"{r.keyword_hits}/{r.keyword_total}" + stop = (r.stop_reason or r.error or "")[:24] + lines.append( + f"| {qid} | {short} | {router} | {r.total_ms} | " + f"{r.llm_steps} | {r.refs} | {kw} | {stop} |\n" + ) + + classic_ms = [r.total_ms for r in runs if r.router == "classic" and not r.error] + agent_ms = [r.total_ms for r in runs if r.router == "agent_explore" and not r.error] + + def _p50(values: list[int]) -> int | None: + if not values: + return None + ordered = sorted(values) + return ordered[len(ordered) // 2] + + lines.extend( + [ + "\n## Aggregate latency (successful runs only)\n", + f"- classic p50: {_p50(classic_ms)} ms ({len(classic_ms)} runs)\n", + f"- agent_explore p50: {_p50(agent_ms)} ms ({len(agent_ms)} runs)\n", + "\n## Notes\n", + "- Runs bypass Redis retrieval cache (direct ``run_retrieval_route``).\n", + "- ``kw hit`` = expected keywords found in ``evidence_text`` " + + "(proxy for recall quality, not a full answer judge).\n", + "- Phase 0 ad-hoc queries are separate; this set is the 8 questions " + + "from ``zh_档案知识库测试样例.docx``.\n", + ] + ) + return "".join(lines) + + +async def main() -> None: + parser = argparse.ArgumentParser(description="Smoke-run agent_explore eval fixture") + parser.add_argument( + "--fixture", + default=str(FIXTURE_PATH), + help="Path to eval queries JSON", + ) + parser.add_argument( + "--output-dir", + "-o", + default=None, + help="Output directory (default: /tmp/agentic_router_eval/)", + ) + parser.add_argument( + "--query-id", + action="append", + default=[], + help="Run only these query ids (repeatable, e.g. --query-id q01)", + ) + parser.add_argument( + "--use-agentic", + choices=["true", "false", "both"], + default="both", + help="classic (false), agent_explore (true), or both", + ) + args = parser.parse_args() + + fixture = json.loads(Path(args.fixture).read_text(encoding="utf-8")) + queries = fixture["queries"] + if args.query_id: + allowed = set(args.query_id) + queries = [q for q in queries if q["id"] in allowed] + + if args.use_agentic == "true": + routers = ["agent_explore"] + elif args.use_agentic == "false": + routers = ["classic"] + else: + routers = list(_MODES) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + output_dir = Path(args.output_dir or f"/tmp/agentic_router_eval/{timestamp}") + output_dir.mkdir(parents=True, exist_ok=True) + + user_id = fixture["corpus"]["user_id"] + namespace = fixture["corpus"]["namespace"] + + all_runs: list[RunMetrics] = [] + for item in queries: + for router in routers: + label = f"{item['id']}_{router}" + print(f"▶ {label}: {item['query'][:60]}…", flush=True) + metrics = await _run_one( + user_id=user_id, + namespace=namespace, + query=item["query"], + query_id=item["id"], + router=router, + use_agentic=router != "classic", + expected_keywords=item.get("expected_keywords") or [], + ) + all_runs.append(metrics) + print( + f" done {metrics.total_ms}ms refs={metrics.refs} " + f"kw={metrics.keyword_hits}/{metrics.keyword_total} " + f"stop={metrics.stop_reason or metrics.error}", + flush=True, + ) + + report = { + "fixture": args.fixture, + "generated_at": datetime.now().isoformat(timespec="seconds"), + "routers": routers, + "runs": [r.__dict__ for r in all_runs], + } + json_path = output_dir / "eval_report.json" + json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") + + md = _render_markdown(fixture, all_runs, output_dir) + md_path = output_dir / "eval_report.md" + md_path.write_text(md, encoding="utf-8") + + print(f"\nWrote {json_path}") + print(f"Wrote {md_path}") + print("\n" + md) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/apps/worker/tests/contract/test_ocr_pages_contract.py b/apps/worker/tests/contract/test_ocr_pages_contract.py index 1fbcde9c..0e0b345a 100644 --- a/apps/worker/tests/contract/test_ocr_pages_contract.py +++ b/apps/worker/tests/contract/test_ocr_pages_contract.py @@ -56,26 +56,29 @@ def test_ocr_pages_writes_joined_text_to_blackboard() -> None: def fake_render(*_args, **_kwargs): return [{"page": 1, "png_path": "/tmp/ocr_page_1.png"}] - def fake_child(worker_fn, page_paths, *, timeout): - assert worker_fn.__name__ == "_run_ocr_worker" - assert page_paths == {1: "/tmp/ocr_page_1.png"} - assert timeout == 300 - return { - "ok": True, - "page_lines": { - 1: [ + calls: list[tuple[int, str, int]] = [] + + class FakeRunner: + def run_page(self, page, image_path, *, timeout): + assert page == 1 + assert image_path == "/tmp/ocr_page_1.png" + calls.append((page, image_path, timeout)) + assert timeout == 1800 + return { + "ok": True, + "page": page, + "lines": [ { "box": [[0, 0], [1, 0], [1, 1], [0, 1]], "text": "Hello", "score": 0.9, } - ] - }, - } + ], + } with ( patch.dict(ocr_pages.__globals__, {"render_pages": fake_render}), - patch.dict(ocr_pages.__globals__, {"run_in_child_process": fake_child}), + patch.dict(ocr_pages.__globals__, {"_OCR_RUNNER": FakeRunner()}), ): result = ocr_pages(ctx, {"pages": [1]}) @@ -85,3 +88,33 @@ def fake_child(worker_fn, page_paths, *, timeout): assert getattr(bands, "header", None) == "" assert getattr(bands, "footer", None) == "" assert result.payload["page_lines"][1][0]["text"] == "Hello" + assert calls == [(1, "/tmp/ocr_page_1.png", 1800)] + + +def test_ocr_pages_runs_one_runner_request_per_page() -> None: + ctx = _ctx() + calls: list[tuple[int, str]] = [] + + def fake_render(*_args: object, **_kwargs: object) -> list[dict[str, object]]: + return [ + {"page": 1, "png_path": "/tmp/ocr_page_1.png"}, + {"page": 2, "png_path": "/tmp/ocr_page_2.png"}, + ] + + class FakeRunner: + def run_page(self, page, image_path, *, timeout): + assert timeout == 1800 + calls.append((page, image_path)) + return {"ok": True, "page": page, "lines": []} + + with ( + patch.dict(ocr_pages.__globals__, {"render_pages": fake_render}), + patch.dict(ocr_pages.__globals__, {"_OCR_RUNNER": FakeRunner()}), + ): + result = ocr_pages(ctx, {"pages": [1, 2]}) + + assert result.status == "ok" + assert calls == [ + (1, "/tmp/ocr_page_1.png"), + (2, "/tmp/ocr_page_2.png"), + ] diff --git a/deploy/ecs/task-definition-api.staging.json b/deploy/ecs/task-definition-api.staging.json index 0b4af2c5..dde385e2 100644 --- a/deploy/ecs/task-definition-api.staging.json +++ b/deploy/ecs/task-definition-api.staging.json @@ -75,7 +75,8 @@ {"name": "LOGFIRE_TOKEN", "valueFrom": "${SECRETS_ARN}:LOGFIRE_TOKEN::"}, {"name": "QSTASH_TOKEN", "valueFrom": "${SECRETS_ARN}:QSTASH_TOKEN::"}, {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, - {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"} + {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"}, + {"name": "CURSOR_API_KEY", "valueFrom": "${SECRETS_ARN}:CURSOR_API_KEY::"} ], "healthCheck": { "command": ["CMD-SHELL", "curl -f http://localhost:5005/health || exit 1"], diff --git a/deploy/ecs/task-definition-worker.staging.json b/deploy/ecs/task-definition-worker.staging.json index 202339d1..8f167696 100644 --- a/deploy/ecs/task-definition-worker.staging.json +++ b/deploy/ecs/task-definition-worker.staging.json @@ -81,7 +81,8 @@ {"name": "LOGFIRE_TOKEN", "valueFrom": "${SECRETS_ARN}:LOGFIRE_TOKEN::"}, {"name": "QSTASH_TOKEN", "valueFrom": "${SECRETS_ARN}:QSTASH_TOKEN::"}, {"name": "QSTASH_CURRENT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_CURRENT_SIGNING_KEY::"}, - {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"} + {"name": "QSTASH_NEXT_SIGNING_KEY", "valueFrom": "${SECRETS_ARN}:QSTASH_NEXT_SIGNING_KEY::"}, + {"name": "CURSOR_API_KEY", "valueFrom": "${SECRETS_ARN}:CURSOR_API_KEY::"} ], "healthCheck": { "command": ["CMD-SHELL", "python -c \"from shared.services.worker_health import assert_worker_healthy; assert_worker_healthy()\""], diff --git a/deprecated/mapnav/README.md b/deprecated/mapnav/README.md new file mode 100644 index 00000000..a7d0c19b --- /dev/null +++ b/deprecated/mapnav/README.md @@ -0,0 +1,26 @@ +# Archived map-nav retrieval route + +This directory holds the retired checklist map-nav episode (PLANNER / HARVEST / CONTROL). + +It is **not imported** by production code. Lint, typecheck, and pytest skip it. + +## Why it was archived + +Retrieval now has two live routes: + +- `use_agentic is False` → classic map-unit BM25 +- otherwise → `agent_explore` (default harness: `cursor_sdk`) + +The `RETRIEVAL_AGENTIC_ROUTER=mapnav` switch was removed. Setting that env var has no effect. + +Shared scoring used by publication and classic recall was extracted first, to `packages/shared-python/shared/services/retrieval/scoring/`. + +## How to restore (manual) + +1. Copy these files back to their original paths: + - `nav/` → `packages/shared-python/shared/services/retrieval/nav/` + - `nav_config.py`, `nav_snapshot.py`, `nav_bridge.py`, `nav_llm_backend.py` → `packages/shared-python/shared/services/retrieval/` + - `trace_mapnav.py` → `packages/shared-python/shared/services/retrieval/trace/mapnav.py` +2. Reverse the Phase 1 import moves: leftover map-nav modules expect `shared.services.retrieval.nav.*` for types that now live under `scoring/`. +3. Re-attach a third branch in `execution/routes.py` (`use_agentic is False` → classic, else map-nav or `agent_explore`). +4. Restore the archived tests from `tests/` and stop excluding this directory from pytest. diff --git a/packages/shared-python/shared/services/retrieval/nav/__init__.py b/deprecated/mapnav/nav/__init__.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/__init__.py rename to deprecated/mapnav/nav/__init__.py diff --git a/packages/shared-python/shared/services/retrieval/nav/_compat.py b/deprecated/mapnav/nav/_compat.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/_compat.py rename to deprecated/mapnav/nav/_compat.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_actions.py b/deprecated/mapnav/nav/nav_actions.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_actions.py rename to deprecated/mapnav/nav/nav_actions.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_address.py b/deprecated/mapnav/nav/nav_address.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_address.py rename to deprecated/mapnav/nav/nav_address.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py b/deprecated/mapnav/nav/nav_agent.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_agent.py rename to deprecated/mapnav/nav/nav_agent.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_assets.py b/deprecated/mapnav/nav/nav_assets.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_assets.py rename to deprecated/mapnav/nav/nav_assets.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_compose.py b/deprecated/mapnav/nav/nav_compose.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_compose.py rename to deprecated/mapnav/nav/nav_compose.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_control.py b/deprecated/mapnav/nav/nav_control.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_control.py rename to deprecated/mapnav/nav/nav_control.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_harvest.py b/deprecated/mapnav/nav/nav_harvest.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_harvest.py rename to deprecated/mapnav/nav/nav_harvest.py diff --git a/deprecated/mapnav/nav/nav_hierarchy.py b/deprecated/mapnav/nav/nav_hierarchy.py new file mode 100644 index 00000000..cabef82e --- /dev/null +++ b/deprecated/mapnav/nav/nav_hierarchy.py @@ -0,0 +1,111 @@ +"""In-memory hierarchy fixtures for map-nav tests. + +``ProviderToolSpace`` / ``NodeMeta`` / ``HierarchyProvider`` live in +``shared.services.retrieval.scoring.hierarchy``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, Set, Tuple + +from shared.services.retrieval.scoring.hierarchy import NodeMeta + +@dataclass +class InMemoryNode: + section_id: str + title: str + content: str = "" + children: List[str] = field(default_factory=list) + + +class InMemoryHierarchyProvider: + """Minimal reference ``HierarchyProvider``: no scoring, no ToolSpace. + + Built directly from a ``{doc_id: [InMemoryNode, ...]}`` map plus a + ``{doc_id: [root_section_id, ...]}`` map — the "hierarchy + summary is + enough" claim's simplest possible witness. + """ + + def __init__( + self, + *, + roots_by_doc: Dict[str, Sequence[str]], + nodes: Dict[str, InMemoryNode], + summaries: Optional[Dict[str, str]] = None, + ) -> None: + self._roots_by_doc = {k: list(v) for k, v in roots_by_doc.items()} + self._nodes = dict(nodes) + self._summaries = dict(summaries or {}) + self._parent: Dict[str, str] = {} + for node in self._nodes.values(): + for child_id in node.children: + self._parent[child_id] = node.section_id + self._owner: Dict[str, str] = {} + for doc_id, root_ids in self._roots_by_doc.items(): + stack = list(root_ids) + while stack: + sid = stack.pop() + if sid in self._owner: + continue + self._owner[sid] = doc_id + node = self._nodes.get(sid) + if node: + stack.extend(node.children) + + def owner_document(self, node_id: str) -> Optional[str]: + return self._owner.get(str(node_id or "").strip()) + + def roots(self, doc_id: str) -> Sequence[str]: + return list(self._roots_by_doc.get(doc_id, ())) + + def children(self, section_id: str) -> Sequence[str]: + node = self._nodes.get(section_id) + return list(node.children) if node else [] + + def node_meta(self, section_id: str) -> NodeMeta: + node = self._nodes.get(section_id) + if node is None: + return NodeMeta() + return NodeMeta( + title=node.title, + summary=self._summaries.get(section_id, ""), + has_children=bool(node.children), + ) + + def parent_id(self, section_id: str) -> Optional[str]: + return self._parent.get(section_id) + + def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: + ancestors: Set[str] = set() + cur = self._parent.get(section_id) + while cur: + ancestors.add(cur) + cur = self._parent.get(cur) + descendants: Set[str] = set() + stack = list(self.children(section_id)) + while stack: + cid = stack.pop() + if cid in descendants: + continue + descendants.add(cid) + stack.extend(self.children(cid)) + return ancestors, descendants + + def content(self, section_id: str) -> str: + node = self._nodes.get(section_id) + if node is None: + return "" + parts: List[str] = [] + + def walk(sid: str) -> None: + cur = self._nodes.get(sid) + if cur is None: + return + if cur.content: + parts.append(cur.content) + for cid in cur.children: + walk(cid) + + walk(section_id) + return "\n".join(parts) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/deprecated/mapnav/nav/nav_knowhere.py similarity index 74% rename from packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py rename to deprecated/mapnav/nav/nav_knowhere.py index cd569631..90d5cf3d 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/deprecated/mapnav/nav/nav_knowhere.py @@ -1,21 +1,8 @@ -"""Knowhere-native hierarchy provider for MAP-NAV. +"""Knowhere-native hierarchy provider leftovers for MAP-NAV. -Knowhere stores everything MAP-NAV needs in two tables: - -``document_sections`` - ``section_id`` (PK) / ``parent_section_id`` (self FK) / ``section_path`` / - ``section_title`` / ``section_level`` / ``summary`` / ``sort_order`` -``document_chunks`` - ``chunk_id`` / ``section_id`` (FK) / ``chunk_type`` / ``content`` / - ``chunk_metadata`` / ``sort_order`` - -``SectionRow`` / ``UnitRow`` mirror those shapes. ``KnowhereProvider`` is a -synchronous in-memory snapshot (so the nav kernel stays sync inside knowhere's -async path). Load from the production Postgres schema via -``load_document_from_db`` / ``load_namespace_from_db`` (local Docker or prod). - -Hierarchy comes from ``parent_section_id`` and depth from ``section_level``, -not from parsing ``section_id`` or ``section_path`` separators. +Publication types and the eager ``KnowhereProvider`` live in +``shared.services.retrieval.scoring``. This module keeps the episode-local +lazy/DB loaders used by map-nav snapshot loading. """ from __future__ import annotations @@ -23,11 +10,9 @@ import os import logging import time -from dataclasses import dataclass, field from hashlib import sha256 from typing import ( Any, - Callable, Dict, Iterable, List, @@ -39,86 +24,25 @@ Tuple, ) -from .nav_address import NavLevel -from .nav_hierarchy import NodeMeta -from .knowhere_hybrid import ( +from shared.services.retrieval.nav.nav_address import NavLevel +from shared.services.retrieval.scoring.hierarchy import NodeMeta +from shared.services.retrieval.scoring.knowhere_hybrid import ( MAP_UNIT_INDEX_FORMAT_VERSION, PersistedScoreCorpus, PersistedScoreUnit, tokenize_query_for_ranker, ) +from shared.services.retrieval.scoring.knowhere_provider import ( + KnowhereProvider, + SectionRow, + UnitRow, + is_root_section_path, +) -_ASSET_TYPES = ("table", "image") -# Knowhere sentinel path for the virtual document container (not a collectable leaf). -ROOT_SECTION_PATH = "Root" _DEFAULT_DSN = "postgresql://root:root123@127.0.0.1:5433/Knowhere" _MAP_SCORE_CHANNELS: Tuple[str, str] = ("path", "content") _logger = logging.getLogger(__name__) - -@dataclass(frozen=True) -class SectionRow: - """One ``document_sections`` row.""" - - section_id: str - parent_section_id: Optional[str] - section_path: str - section_title: str - section_level: int - summary: str - sort_order: int - - -@dataclass(frozen=True) -class UnitRow: - """One ``document_chunks`` row.""" - - chunk_id: str - section_id: Optional[str] - chunk_type: str - content: str - sort_order: int - source_chunk_path: str = "" - file_path: str = "" - metadata: Dict[str, Any] = field(default_factory=dict) - - -def asset_display_text(unit: UnitRow) -> str: - """Body text for an asset unit, whose ``content`` is only a file path. - - Mirrors knowhere's own assembly: an asset contributes its summary, not its - path. Without this an asset unit is unscorable and unreadable. - """ - meta = unit.metadata or {} - title = str(meta.get("asset_title") or "").strip() - summary = str(meta.get("summary") or "").strip() - ref = unit.file_path or unit.source_chunk_path or unit.content - label = "Table" if unit.chunk_type == "table" else "Image" - parts = [f"[{label}: {ref}]"] if ref else [f"[{label}]"] - if title: - parts.append(title) - if summary: - parts.append(summary) - return "\n".join(parts) - - -def normalize_section_path(path: str) -> str: - """Canonical path for gold/lookup: ``a / b`` (accepts ``a/b`` or ``a / b``).""" - raw = str(path or "").strip().strip("/") - if not raw or raw == ROOT_SECTION_PATH: - return "" - if " / " in raw: - parts = [p.strip() for p in raw.split(" / ") if p.strip()] - else: - parts = [p.strip() for p in raw.split("/") if p.strip()] - return " / ".join(parts) - - -def is_root_section_path(path: str) -> bool: - """True when the raw ``section_path`` is Knowhere's Root container.""" - return str(path or "").strip() == ROOT_SECTION_PATH - - def is_root_section(provider_or_ts: Any, section_id: str) -> bool: """True when ``section_id`` is a Root container (raw path, not normalized).""" sid = str(section_id or "").strip() @@ -134,21 +58,6 @@ def is_root_section(provider_or_ts: Any, section_id: str) -> bool: return False -def _connect_to_targets(metadata: Dict[str, Any]) -> List[str]: - """``chunk_metadata.connect_to[].target`` ids (document order, first wins upstream).""" - raw = metadata.get("connect_to") if isinstance(metadata, dict) else None - if not isinstance(raw, list): - return [] - out: List[str] = [] - for conn in raw: - if not isinstance(conn, dict): - continue - target = str(conn.get("target") or "").strip() - if target: - out.append(target) - return out - - def knowhere_database_url() -> str: configured = ( str(os.environ.get("KNOWHERE_DATABASE_URL") or "").strip() @@ -300,7 +209,7 @@ def load_persisted_score_corpus( ``document_map_unit_tokens`` filtered to the query tokens. Average IDF comes from ``document_map_unit_indexes`` (written at index time). """ - from shared.services.retrieval.nav.persisted_score_load import ( + from shared.services.retrieval.scoring.persisted_score_load import ( build_channel_bm25_stats, combine_average_idf, ) @@ -646,252 +555,6 @@ def _unit_from_row(row: Sequence[object]) -> UnitRow: ) -class KnowhereProvider: - """``HierarchyProvider`` over knowhere section/chunk rows.""" - - def __init__( - self, - *, - doc_id: str, - sections: Sequence[SectionRow], - units: Sequence[UnitRow], - lazy_loader: Optional[Callable[[str], Sequence[UnitRow]]] = None, - known_chunk_ids: Optional[Sequence[str]] = None, - ) -> None: - self.doc_id = str(doc_id) - self._lazy_loader = lazy_loader - self._loaded_sections: Set[str] = set() - self._sections: Dict[str, SectionRow] = {s.section_id: s for s in sections} - self._children: Dict[str, List[str]] = {} - self._roots: List[str] = [] - self._path_to_id: Dict[str, str] = {} - for row in sorted(sections, key=lambda s: (s.sort_order, s.section_id)): - parent = row.parent_section_id - if parent and parent in self._sections: - self._children.setdefault(parent, []).append(row.section_id) - else: - self._roots.append(row.section_id) - key = normalize_section_path(row.section_path) - if key: - self._path_to_id[key] = row.section_id - - self._units_by_section: Dict[str, List[UnitRow]] = {} - self._chunk_ids: Set[str] = set() - if known_chunk_ids: - self._chunk_ids.update( - str(chunk_id).strip() - for chunk_id in known_chunk_ids - if str(chunk_id).strip() - ) - for unit in sorted(units, key=lambda u: (u.sort_order, u.chunk_id)): - sid = unit.section_id - if not sid or sid not in self._sections: - continue - self._units_by_section.setdefault(sid, []).append(unit) - if unit.chunk_id: - self._chunk_ids.add(unit.chunk_id) - self._remount_root_assets() - - def _ensure_section_loaded(self, section_id: str) -> None: - if self._lazy_loader is None or section_id in self._loaded_sections: - return - loaded = list(self._lazy_loader(section_id) or ()) - self._loaded_sections.add(section_id) - if not loaded: - return - current = self._units_by_section.setdefault(section_id, []) - known = {unit.chunk_id for unit in current} - for unit in loaded: - if unit.chunk_id and unit.chunk_id not in known: - current.append(unit) - known.add(unit.chunk_id) - current.sort(key=lambda unit: (unit.sort_order, unit.chunk_id)) - - def _remount_root_assets(self) -> None: - """Reattach Root-FK image|table units to host sections via ``connect_to``. - - Aligns with Knowhere ``resolve_root_asset_owners``: assets whose FK still - points at Root are owned by the text chunk that lists them in - ``metadata.connect_to``. Unresolved Root assets leave the evidence surface. - """ - root_sids = [ - sid - for sid, row in self._sections.items() - if is_root_section_path(row.section_path) - ] - if not root_sids: - return - - root_assets: Dict[str, UnitRow] = {} - for sid in root_sids: - for unit in self._units_by_section.get(sid, ()): - if unit.chunk_type in _ASSET_TYPES and unit.chunk_id: - root_assets[unit.chunk_id] = unit - if not root_assets: - return - - owner_by_asset: Dict[str, str] = {} - for sid, units in self._units_by_section.items(): - row = self._sections.get(sid) - if row is None or is_root_section_path(row.section_path): - continue - for unit in units: - if unit.chunk_type != "text": - continue - for target in _connect_to_targets(unit.metadata or {}): - if target in root_assets and target not in owner_by_asset: - owner_by_asset[target] = sid - - touched_owners: Set[str] = set() - for chunk_id, owner_sid in owner_by_asset.items(): - unit = root_assets[chunk_id] - remounted = UnitRow( - chunk_id=unit.chunk_id, - section_id=owner_sid, - chunk_type=unit.chunk_type, - content=unit.content, - sort_order=unit.sort_order, - source_chunk_path=unit.source_chunk_path, - file_path=unit.file_path, - metadata=dict(unit.metadata or {}), - ) - self._units_by_section.setdefault(owner_sid, []).append(remounted) - touched_owners.add(owner_sid) - - for sid in root_sids: - self._units_by_section[sid] = [ - u - for u in self._units_by_section.get(sid, ()) - if u.chunk_type not in _ASSET_TYPES - ] - for sid in touched_owners: - self._units_by_section[sid].sort(key=lambda u: (u.sort_order, u.chunk_id)) - - def address_level(self, node_id: str) -> Optional[NavLevel]: - sid = str(node_id or "").strip() - if not sid: - return NavLevel.NAMESPACE - if sid == self.doc_id: - return NavLevel.DOCUMENT - if sid in self._sections: - return NavLevel.SECTION - if sid in self._chunk_ids: - return NavLevel.CHUNK - return None - - def owner_document(self, node_id: str) -> Optional[str]: - sid = str(node_id or "").strip() - if not sid: - return None - if sid == self.doc_id or sid in self._sections or sid in self._chunk_ids: - return self.doc_id - return None - - def roots(self, doc_id: str) -> Sequence[str]: - return list(self._roots) if str(doc_id) == self.doc_id else [] - - def children(self, section_id: str) -> Sequence[str]: - return list(self._children.get(section_id, ())) - - def node_meta(self, section_id: str) -> NodeMeta: - row = self._sections.get(section_id) - if row is None: - return NodeMeta() - return NodeMeta( - title=row.section_title, - summary=row.summary, - has_children=bool(self._children.get(section_id)), - ) - - def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: - ancestors: Set[str] = set() - cur = self._sections.get(section_id) - while cur is not None and cur.parent_section_id: - parent = cur.parent_section_id - if parent in ancestors: - break - ancestors.add(parent) - cur = self._sections.get(parent) - descendants: Set[str] = set() - stack = list(self.children(section_id)) - while stack: - cid = stack.pop() - if cid in descendants: - continue - descendants.add(cid) - stack.extend(self.children(cid)) - return ancestors, descendants - - def content(self, section_id: str) -> str: - units = self.subtree_units(section_id) - return "\n".join(self.unit_text(u) for u in units if self.unit_text(u)) - - def self_units(self, section_id: str) -> List[UnitRow]: - self._ensure_section_loaded(section_id) - return list(self._units_by_section.get(section_id, ())) - - def subtree_units(self, section_id: str) -> List[UnitRow]: - out = list(self.self_units(section_id)) - for cid in self.relations(section_id)[1]: - out.extend(self.self_units(cid)) - out.sort(key=lambda u: (u.sort_order, u.chunk_id)) - return out - - def leaf_ids(self, section_id: str) -> List[str]: - out: List[str] = [] - - def rec(sid: str) -> None: - kids = self.children(sid) - if not kids: - out.append(sid) - return - for kid in kids: - rec(kid) - - rec(section_id) - return out - - def path_titles(self, section_id: str) -> str: - chain: List[str] = [] - cur = self._sections.get(section_id) - while cur is not None: - if cur.section_title: - chain.append(cur.section_title) - parent = cur.parent_section_id - cur = self._sections.get(parent) if parent else None - return " / ".join(reversed(chain)) - - def parent_id(self, section_id: str) -> Optional[str]: - row = self._sections.get(section_id) - return row.parent_section_id if row else None - - def section_path(self, section_id: str) -> str: - row = self._sections.get(section_id) - return str(row.section_path or "") if row else "" - - def resolve_path(self, path: str) -> Optional[str]: - """Map a human/gold path to ``section_id`` (``sec_*``).""" - key = normalize_section_path(path) - if not key: - return None - return self._path_to_id.get(key) - - def unit_text(self, unit: UnitRow) -> str: - if unit.chunk_type in _ASSET_TYPES: - return asset_display_text(unit) - return str(unit.content or "").strip() - - def summaries(self) -> Dict[str, str]: - return { - sid: row.summary - for sid, row in self._sections.items() - if str(row.summary or "").strip() - } - - def all_section_ids(self) -> List[str]: - return list(self._sections) - - class LazyKnowhereProvider(KnowhereProvider): """Hierarchy provider that loads full chunk rows only on first access.""" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py b/deprecated/mapnav/nav/nav_llm.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_llm.py rename to deprecated/mapnav/nav/nav_llm.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/deprecated/mapnav/nav/nav_map_scores.py similarity index 52% rename from packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py rename to deprecated/mapnav/nav/nav_map_scores.py index 35ffd890..8877451a 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/deprecated/mapnav/nav/nav_map_scores.py @@ -4,18 +4,8 @@ import time from typing import Any, Dict, List, Optional, Sequence, Set, Tuple -from .knowhere_hybrid import ( - build_content_search_text, - build_path_search_text, - build_term_search_text, - PersistedScoreCorpus, - PersistedScoreUnit, - score_persisted_corpus_many, -) -from .persisted_score_load import ( - average_idf_from_unit_dfs, - build_channel_bm25_stats, -) +from shared.services.retrieval.scoring.knowhere_hybrid import score_persisted_corpus_many +from shared.services.retrieval.scoring.score_units import _walk_tree _logger = logging.getLogger(__name__) @@ -35,188 +25,6 @@ def _count_tree_shape( leaf_sections: int = sum(len(value[1]) for value in tree_by_doc.values()) return section_nodes, section_edges, leaf_sections - -def _build_legacy_score_corpus(ts: Any, doc_ids: Sequence[str]) -> PersistedScoreCorpus: - """Build the retired in-memory scorer input when persisted indexes are absent.""" - raw_units: List[dict] = [] - for doc_id in doc_ids: - raw_units.extend(build_score_units(ts, doc_id)) - frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} - unit_rows: List[dict] = [] - path_dfs: Dict[str, int] = {} - content_dfs: Dict[str, int] = {} - for unit in raw_units: - unit_id = str(unit.get("chunk_id") or "").strip() - if not unit_id: - continue - path_tokens = str(unit.get("path_search_text") or "").split() - content_tokens = str(unit.get("content_search_text") or "").split() - path_freq: Dict[str, int] = {} - content_freq: Dict[str, int] = {} - for token in path_tokens: - path_freq[token] = path_freq.get(token, 0) + 1 - for token in content_tokens: - content_freq[token] = content_freq.get(token, 0) + 1 - frequencies[(unit_id, "path")] = path_freq - frequencies[(unit_id, "content")] = content_freq - for token in path_freq: - path_dfs[token] = path_dfs.get(token, 0) + 1 - for token in content_freq: - content_dfs[token] = content_dfs.get(token, 0) + 1 - unit_rows.append( - { - "unit_id": unit_id, - "path_length": len(path_tokens), - "content_length": len(content_tokens), - } - ) - unit_count = len(unit_rows) - return PersistedScoreCorpus( - units=[ - PersistedScoreUnit( - unit_id=str(row["unit_id"]), - path_length=int(row["path_length"]), - content_length=int(row["content_length"]), - path_frequencies=frequencies[(str(row["unit_id"]), "path")], - content_frequencies=frequencies[(str(row["unit_id"]), "content")], - ) - for row in unit_rows - ], - path_stats=build_channel_bm25_stats( - unit_rows=unit_rows, - map_unit_id_field="unit_id", - length_field="path_length", - channel="path", - query_tokens=list(path_dfs), - frequencies=frequencies, - average_idf=average_idf_from_unit_dfs( - unit_count=unit_count, token_document_frequency=path_dfs - ), - ), - content_stats=build_channel_bm25_stats( - unit_rows=unit_rows, - map_unit_id_field="unit_id", - length_field="content_length", - channel="content", - query_tokens=list(content_dfs), - frequencies=frequencies, - average_idf=average_idf_from_unit_dfs( - unit_count=unit_count, token_document_frequency=content_dfs - ), - ), - ) - - -def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: - children_fn = getattr(ts, "_children_for_section_path", None) - if not callable(children_fn): - st = ts.get_structure(section_id) - rows = st.get("children") or [] - return [ - str(r.get("section_id") or "").strip() for r in rows if r.get("section_id") - ] - rows = children_fn(section_id, doc_id) - return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] - - -def _line_content(ts: Any, section_id: str, doc_id: str) -> str: - """Raw line text for a section node (no truncation).""" - idx = getattr(ts, "_idx", None) - b = getattr(idx, "_bundles", {}).get(doc_id) if idx is not None else None - if b is None: - path_fn = getattr(ts, "path_titles", None) - if callable(path_fn): - path = str(path_fn(section_id, doc_id) or "").strip() - return path.rsplit(" / ", 1)[-1] if path else "" - st = ts.get_structure(section_id) - return str(st.get("preview") or "").strip() - loc = getattr(idx, "_node_to_doc_line", {}).get(section_id) - if not loc: - return "" - _doc, line_idx = loc - if line_idx < 0 or line_idx >= len(b.lines): - return "" - return str(b.lines[line_idx].content or "").strip() - - -def _ancestor_path_titles(ts: Any, section_id: str, doc_id: str) -> str: - idx = getattr(ts, "_idx", None) - if idx is None: - # Provider-backed spaces expose the title chain directly; without this - # the path channel would score every unit as empty. - path_fn = getattr(ts, "path_titles", None) - return str(path_fn(section_id, doc_id) or "") if callable(path_fn) else "" - try: - ancestors = list(idx.ancestor_line_node_ids(section_id)) - except Exception: - ancestors = [] - titles: List[str] = [] - for aid in reversed(ancestors): - if not str(aid).startswith(f"{doc_id}:"): - continue - titles.append(_line_content(ts, aid, doc_id)) - titles.append(_line_content(ts, section_id, doc_id)) - return " / ".join(t for t in titles if t) - - -def _self_only_text(ts: Any, section_id: str, doc_id: str) -> Tuple[str, bool]: - """Return (self_text, has_interstitial_body). - - Interstitial means self_only span contains content beyond the heading line - itself (structural: more than one line/chunk in the self span). - """ - self_fn = getattr(ts, "materialize_self_only_chunks", None) - if not callable(self_fn): - return "", False - chunks = list(self_fn(section_id, doc_id) or []) - if not chunks: - return "", False - texts = [str(getattr(c, "text", "") or "").strip() for c in chunks] - texts = [t for t in texts if t] - if not texts: - return "", False - # Structural interstitial: self span covers more than the node heading line. - has_interstitial = len(chunks) > 1 - return "\n".join(texts), has_interstitial - - -def _section_body_text(ts: Any, section_id: str, doc_id: str) -> str: - """Heading + lines until first structural child (leaf body / parent self span).""" - text, _ = _self_only_text(ts, section_id, doc_id) - if text: - return text - return _line_content(ts, section_id, doc_id) - - -def _walk_tree( - ts: Any, - doc_id: str, - root_ids: Sequence[str], -) -> Tuple[Dict[str, List[str]], Set[str], Dict[str, str]]: - """Return children map, leaf ids, and title map for reachable nodes.""" - children_map: Dict[str, List[str]] = {} - titles: Dict[str, str] = {} - leaves: Set[str] = set() - seen: Set[str] = set() - - def walk(sid: str) -> None: - if not sid or sid in seen: - return - seen.add(sid) - titles[sid] = _line_content(ts, sid, doc_id) - kids = [c for c in _children_ids(ts, sid, doc_id) if c] - children_map[sid] = kids - if not kids: - leaves.add(sid) - return - for kid in kids: - walk(kid) - - for rid in root_ids: - walk(rid) - return children_map, leaves, titles - - def _collect_descendant_leaves( section_id: str, children_map: Dict[str, List[str]], @@ -271,74 +79,6 @@ def _pool_unit_scores_to_tree( return map_scores -def build_score_units( - ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None -) -> List[dict]: - """Build leaf (+ interstitial self_only) units for hybrid scoring.""" - if root_ids is None: - root_ids = list(ts.sections_for_doc(doc_id)) - children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) - units: List[dict] = [] - seen_unit_ids: Set[str] = set() - - for leaf_id in sorted(leaves): - content = _section_body_text(ts, leaf_id, doc_id) or ( - titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) - ) - path_text = _ancestor_path_titles(ts, leaf_id, doc_id) - unit_id = leaf_id - if unit_id in seen_unit_ids: - continue - seen_unit_ids.add(unit_id) - title = titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) - units.append( - { - "chunk_id": unit_id, - "section_id": leaf_id, - "kind": "leaf", - "content": content, - "path_text": path_text, - "path_search_text": build_path_search_text( - section_path=path_text, section_title=title or content - ), - "content_search_text": build_content_search_text(content), - "term_search_text": build_term_search_text( - content, path_text=path_text - ), - } - ) - - # Parents with interstitial self body. - for sid, kids in children_map.items(): - if not kids: - continue - self_text, has_interstitial = _self_only_text(ts, sid, doc_id) - if not has_interstitial or not self_text: - continue - unit_id = f"{sid}__self" - if unit_id in seen_unit_ids: - continue - seen_unit_ids.add(unit_id) - path_text = _ancestor_path_titles(ts, sid, doc_id) - units.append( - { - "chunk_id": unit_id, - "section_id": sid, - "kind": "self_only", - "content": self_text, - "path_text": path_text, - "path_search_text": build_path_search_text( - section_path=path_text, section_title=titles.get(sid) or "" - ), - "content_search_text": build_content_search_text(self_text), - "term_search_text": build_term_search_text( - self_text, path_text=path_text - ), - } - ) - return units - - def compute_map_scores( ts: Any, *, @@ -453,13 +193,6 @@ def compute_corpus_map_and_unit_scores_many( time.perf_counter() - loader_started, persisted_corpus is not None, ) - if persisted_corpus is None: - _logger.warning( - "retrieval map index unavailable; using bounded legacy in-memory scorer " - "documents=%d", - len(valid_doc_ids), - ) - persisted_corpus = _build_legacy_score_corpus(ts, valid_doc_ids) score_started = time.perf_counter() unit_scores_by_query = ( score_persisted_corpus_many(persisted_corpus, unique_queries) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_navigate.py b/deprecated/mapnav/nav/nav_navigate.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_navigate.py rename to deprecated/mapnav/nav/nav_navigate.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_node_filter.py b/deprecated/mapnav/nav/nav_node_filter.py similarity index 70% rename from packages/shared-python/shared/services/retrieval/nav/nav_node_filter.py rename to deprecated/mapnav/nav/nav_node_filter.py index 45b2c1eb..5e895ee2 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_node_filter.py +++ b/deprecated/mapnav/nav/nav_node_filter.py @@ -7,21 +7,14 @@ from __future__ import annotations -import re from dataclasses import dataclass, field -from typing import Any, Dict, Iterable, List, Literal, Sequence, Tuple +from typing import Any, Dict, Iterable, List, Sequence, Tuple -MatchKind = Literal["substring", "regex"] -FilterField = Literal["path", "summary"] - -_MAX_REGEX_PATTERN_LEN = 256 - - -@dataclass(frozen=True) -class FieldPredicate: - field: FilterField - terms: Tuple[str, ...] - match: MatchKind = "substring" +from shared.services.retrieval.scoring.node_filter_predicates import ( + FieldPredicate, + _compile_predicates, + _node_matches, +) @dataclass(frozen=True) @@ -36,26 +29,9 @@ class FilterResult: cardinality: int failed_predicates: List[str] = field(default_factory=list) - -def field_predicate( - field: str, - terms: Sequence[str], - match: str = "substring", -) -> FieldPredicate: - key = str(field or "").strip().lower() - if key not in {"path", "summary"}: - raise ValueError(f"unsupported filter field: {field!r}") - kind = str(match or "substring").strip().lower() - if kind not in {"substring", "regex"}: - raise ValueError(f"unsupported filter match: {match!r}") - cleaned = tuple(str(term) for term in terms if str(term)) - return FieldPredicate(field=key, terms=cleaned, match=kind) # type: ignore[arg-type] - - def node_filter(predicates: Sequence[FieldPredicate]) -> NodeFilter: return NodeFilter(predicates=tuple(predicates)) - def apply_node_filter( ts: Any, doc_ids: Sequence[str], @@ -133,52 +109,6 @@ def render_submap_observation( lines.append("\n".join(block)) return "\n".join(lines) - -def _compile_predicates( - predicates: Sequence[FieldPredicate], -) -> Tuple[List[Tuple[FieldPredicate, List[Any]]], List[str]]: - compiled: List[Tuple[FieldPredicate, List[Any]]] = [] - failed: List[str] = [] - for pred in predicates: - if pred.match != "regex": - compiled.append((pred, [])) - continue - patterns: List[Any] = [] - ok = True - for term in pred.terms: - if len(term) > _MAX_REGEX_PATTERN_LEN: - failed.append(f"{pred.field}:regex:too_long") - ok = False - break - try: - patterns.append(re.compile(term, flags=re.IGNORECASE)) - except re.error: - failed.append(f"{pred.field}:regex:invalid") - ok = False - break - if ok: - compiled.append((pred, patterns)) - return compiled, failed - - -def _node_matches( - values: Dict[str, str], - compiled: Sequence[Tuple[FieldPredicate, List[Any]]], -) -> bool: - if not compiled: - return True - for pred, patterns in compiled: - text = values.get(pred.field, "") - if pred.match == "regex": - if not patterns or not any(p.search(text or "") for p in patterns): - return False - continue - haystack = (text or "").lower() - if not pred.terms or not any(term.lower() in haystack for term in pred.terms): - return False - return True - - def _iter_doc_nodes(ts: Any, doc_id: str) -> Iterable[Tuple[str, str, bool]]: yield doc_id, doc_id, True stack = list(_roots(ts, doc_id)) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/deprecated/mapnav/nav/nav_orchestrate.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py rename to deprecated/mapnav/nav/nav_orchestrate.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_plan.py b/deprecated/mapnav/nav/nav_plan.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_plan.py rename to deprecated/mapnav/nav/nav_plan.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_policy.py b/deprecated/mapnav/nav/nav_policy.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_policy.py rename to deprecated/mapnav/nav/nav_policy.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_projection.py b/deprecated/mapnav/nav/nav_projection.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_projection.py rename to deprecated/mapnav/nav/nav_projection.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_scope_filter.py b/deprecated/mapnav/nav/nav_scope_filter.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/nav/nav_scope_filter.py rename to deprecated/mapnav/nav/nav_scope_filter.py index 886a0527..a381a9a8 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_scope_filter.py +++ b/deprecated/mapnav/nav/nav_scope_filter.py @@ -11,11 +11,12 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Sequence +from shared.services.retrieval.scoring.node_filter_predicates import field_predicate + from .nav_node_filter import ( FilterResult, NodeFilter, apply_node_filter, - field_predicate, node_filter, render_submap_observation, ) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_token_budget.py b/deprecated/mapnav/nav/nav_token_budget.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_token_budget.py rename to deprecated/mapnav/nav/nav_token_budget.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_types.py b/deprecated/mapnav/nav/nav_types.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_types.py rename to deprecated/mapnav/nav/nav_types.py diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_verify.py b/deprecated/mapnav/nav/nav_verify.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav/nav_verify.py rename to deprecated/mapnav/nav/nav_verify.py diff --git a/packages/shared-python/shared/services/retrieval/nav_bridge.py b/deprecated/mapnav/nav_bridge.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav_bridge.py rename to deprecated/mapnav/nav_bridge.py diff --git a/packages/shared-python/shared/services/retrieval/nav_config.py b/deprecated/mapnav/nav_config.py similarity index 94% rename from packages/shared-python/shared/services/retrieval/nav_config.py rename to deprecated/mapnav/nav_config.py index 48067194..c7338f09 100644 --- a/packages/shared-python/shared/services/retrieval/nav_config.py +++ b/deprecated/mapnav/nav_config.py @@ -16,10 +16,11 @@ from typing import Any from shared.services.retrieval.nav.nav_types import NavConfig +from shared.services.retrieval.settings import EVIDENCE_TEXT_CHAR_BUDGET # Migrated probe / llm_api.env stack. MAPNAV_MODEL = "deepseek-v4-flash" -MAPNAV_EVIDENCE_CHARS = 12_000 +MAPNAV_EVIDENCE_CHARS = EVIDENCE_TEXT_CHAR_BUDGET MAPNAV_TOKEN_LIMIT = 100_000 MAPNAV_PLANNER_THINK_MAX_TOKENS = 16_384 MAPNAV_TRACE_RAW_CHARS = 2_000 diff --git a/packages/shared-python/shared/services/retrieval/nav_llm_backend.py b/deprecated/mapnav/nav_llm_backend.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/nav_llm_backend.py rename to deprecated/mapnav/nav_llm_backend.py diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/deprecated/mapnav/nav_snapshot.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/nav_snapshot.py rename to deprecated/mapnav/nav_snapshot.py index 9f92617c..5e1016dd 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/deprecated/mapnav/nav_snapshot.py @@ -41,12 +41,14 @@ from shared.models.database.job_result import JobResult from shared.services.retrieval.nav.nav_knowhere import ( LazyKnowhereProvider, - KnowhereProvider, NamespaceKnowhereProvider, ReadOnlyChunkStore, + knowhere_database_url, +) +from shared.services.retrieval.scoring.knowhere_provider import ( + KnowhereProvider, SectionRow, UnitRow, - knowhere_database_url, ) from shared.services.retrieval.search.section_filters import is_excluded_section from shared.services.retrieval.serving_manifest import ( diff --git a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_snapshot_quality_contract.py similarity index 97% rename from apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_lazy_snapshot_quality_contract.py index c796e2a1..00855e6b 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py +++ b/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_snapshot_quality_contract.py @@ -6,26 +6,28 @@ import math from typing import Any -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace from shared.services.retrieval.nav.nav_knowhere import ( - KnowhereProvider, LazyKnowhereProvider, NamespaceKnowhereProvider, - SectionRow, - UnitRow, knowhere_database_url, ) from shared.services.retrieval.nav.nav_map_scores import ( - build_score_units, compute_corpus_map_and_unit_scores, compute_corpus_map_and_unit_scores_many, ) -from shared.services.retrieval.nav.knowhere_hybrid import ( +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_hybrid import ( PersistedBm25Stats, PersistedScoreCorpus, PersistedScoreUnit, score_persisted_corpus_many, ) +from shared.services.retrieval.scoring.knowhere_provider import ( + KnowhereProvider, + SectionRow, + UnitRow, +) +from shared.services.retrieval.scoring.score_units import build_score_units @dataclass diff --git a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_tree_contract.py similarity index 93% rename from apps/api/tests/contract/test_retrieval_lazy_tree_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_lazy_tree_contract.py index 9571f151..038756a9 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py +++ b/deprecated/mapnav/tests/api-contract/test_retrieval_lazy_tree_contract.py @@ -2,9 +2,9 @@ from __future__ import annotations -from shared.services.retrieval.nav.nav_hierarchy import NodeMeta, ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import KnowhereProvider, SectionRow -from shared.services.retrieval.nav.nav_map_scores import _walk_tree +from shared.services.retrieval.scoring.hierarchy import NodeMeta, ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import KnowhereProvider, SectionRow +from shared.services.retrieval.scoring.score_units import _walk_tree from shared.services.retrieval.nav._compat import Chunk from shared.services.retrieval.nav.nav_compose import pack_nav_evidence from shared.services.retrieval.nav.nav_types import NavConfig, NavState diff --git a/apps/api/tests/contract/test_retrieval_map_score_parity_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_map_score_parity_contract.py similarity index 100% rename from apps/api/tests/contract/test_retrieval_map_score_parity_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_map_score_parity_contract.py diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_map_unit_index_contract.py similarity index 99% rename from apps/api/tests/contract/test_retrieval_map_unit_index_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_map_unit_index_contract.py index 5128382f..9cde1b65 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/deprecated/mapnav/tests/api-contract/test_retrieval_map_unit_index_contract.py @@ -13,23 +13,25 @@ DocumentMapUnitIndex, DocumentMapUnitToken, ) -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace from shared.services.retrieval.nav._compat import Chunk, EpisodeResult from shared.services.retrieval.nav import nav_knowhere +from shared.services.retrieval.nav.nav_knowhere import ( + LazyKnowhereProvider, + NamespaceKnowhereProvider, + ReadOnlyChunkStore, +) from shared.services.retrieval.nav.nav_map_scores import ( - build_score_units, compute_corpus_map_and_unit_scores, select_map_highlights, ) from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows -from shared.services.retrieval.nav.nav_knowhere import ( +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, - LazyKnowhereProvider, - NamespaceKnowhereProvider, - ReadOnlyChunkStore, SectionRow, UnitRow, ) +from shared.services.retrieval.scoring.score_units import build_score_units from shared.services.retrieval.nav_snapshot import load_nav_snapshot from shared.services.retrieval.publication_content import ( replace_document_revision_content, diff --git a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py similarity index 98% rename from apps/api/tests/contract/test_retrieval_mapnav_session_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py index d267941b..f5ca83f6 100644 --- a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py +++ b/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py @@ -16,7 +16,7 @@ ) 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.scoring.knowhere_provider import SectionRow, UnitRow from shared.services.retrieval.nav_snapshot import build_nav_snapshot RouteRow = dict[str, object] diff --git a/apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_relit_map_cache_contract.py similarity index 100% rename from apps/api/tests/contract/test_retrieval_relit_map_cache_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_relit_map_cache_contract.py diff --git a/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_batching_contract.py similarity index 100% rename from apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_batching_contract.py diff --git a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_consistency_contract.py similarity index 100% rename from apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_consistency_contract.py diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_large_corpus_contract.py similarity index 100% rename from apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_large_corpus_contract.py diff --git a/apps/api/tests/contract/test_retrieval_snapshot_redis_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_redis_contract.py similarity index 100% rename from apps/api/tests/contract/test_retrieval_snapshot_redis_contract.py rename to deprecated/mapnav/tests/api-contract/test_retrieval_snapshot_redis_contract.py diff --git a/packages/shared-python/shared/tests/test_nav_bridge_config.py b/deprecated/mapnav/tests/shared/test_nav_bridge_config.py similarity index 98% rename from packages/shared-python/shared/tests/test_nav_bridge_config.py rename to deprecated/mapnav/tests/shared/test_nav_bridge_config.py index 89e8469d..1ee16569 100644 --- a/packages/shared-python/shared/tests/test_nav_bridge_config.py +++ b/deprecated/mapnav/tests/shared/test_nav_bridge_config.py @@ -14,7 +14,7 @@ os.environ.setdefault("S3_TEMP_PATH", "/tmp") from shared.services.retrieval.nav._compat import Chunk -from shared.services.retrieval.nav.nav_knowhere import SectionRow, UnitRow +from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow from shared.services.retrieval.nav_bridge import build_referenced_chunks from shared.services.retrieval.nav_config import ( MAPNAV_MODEL, diff --git a/packages/shared-python/shared/tests/test_nav_llm_backend.py b/deprecated/mapnav/tests/shared/test_nav_llm_backend.py similarity index 100% rename from packages/shared-python/shared/tests/test_nav_llm_backend.py rename to deprecated/mapnav/tests/shared/test_nav_llm_backend.py diff --git a/packages/shared-python/shared/tests/test_nav_node_filter.py b/deprecated/mapnav/tests/shared/test_nav_node_filter.py similarity index 93% rename from packages/shared-python/shared/tests/test_nav_node_filter.py rename to deprecated/mapnav/tests/shared/test_nav_node_filter.py index ef662b64..bca04462 100644 --- a/packages/shared-python/shared/tests/test_nav_node_filter.py +++ b/deprecated/mapnav/tests/shared/test_nav_node_filter.py @@ -2,18 +2,18 @@ from __future__ import annotations -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( +from shared.services.retrieval.nav.nav_knowhere import NamespaceKnowhereProvider +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, - NamespaceKnowhereProvider, SectionRow, ) from shared.services.retrieval.nav.nav_node_filter import ( apply_node_filter, - field_predicate, node_filter, render_submap_observation, ) +from shared.services.retrieval.scoring.node_filter_predicates import field_predicate def _section( diff --git a/packages/shared-python/shared/tests/test_nav_node_filter_wire.py b/deprecated/mapnav/tests/shared/test_nav_node_filter_wire.py similarity index 97% rename from packages/shared-python/shared/tests/test_nav_node_filter_wire.py rename to deprecated/mapnav/tests/shared/test_nav_node_filter_wire.py index f3129736..1025a647 100644 --- a/packages/shared-python/shared/tests/test_nav_node_filter_wire.py +++ b/deprecated/mapnav/tests/shared/test_nav_node_filter_wire.py @@ -4,10 +4,10 @@ from typing import Any -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( +from shared.services.retrieval.nav.nav_knowhere import NamespaceKnowhereProvider +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, - NamespaceKnowhereProvider, SectionRow, ) from shared.services.retrieval.nav.nav_orchestrate import _execute_subgoal_harvest_once diff --git a/packages/shared-python/shared/tests/test_nav_plan_node_filter.py b/deprecated/mapnav/tests/shared/test_nav_plan_node_filter.py similarity index 100% rename from packages/shared-python/shared/tests/test_nav_plan_node_filter.py rename to deprecated/mapnav/tests/shared/test_nav_plan_node_filter.py diff --git a/packages/shared-python/shared/tests/test_nav_plan_query_only.py b/deprecated/mapnav/tests/shared/test_nav_plan_query_only.py similarity index 100% rename from packages/shared-python/shared/tests/test_nav_plan_query_only.py rename to deprecated/mapnav/tests/shared/test_nav_plan_query_only.py diff --git a/packages/shared-python/shared/tests/test_nav_projection_prod.py b/deprecated/mapnav/tests/shared/test_nav_projection_prod.py similarity index 95% rename from packages/shared-python/shared/tests/test_nav_projection_prod.py rename to deprecated/mapnav/tests/shared/test_nav_projection_prod.py index 1d5219d6..d03a30ba 100644 --- a/packages/shared-python/shared/tests/test_nav_projection_prod.py +++ b/deprecated/mapnav/tests/shared/test_nav_projection_prod.py @@ -12,8 +12,8 @@ os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") os.environ.setdefault("S3_TEMP_PATH", "/tmp") -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import SectionRow, UnitRow +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow from shared.services.retrieval.nav.nav_projection import ( _section_summary_for_map, build_map, diff --git a/packages/shared-python/shared/tests/test_nav_scope_filter.py b/deprecated/mapnav/tests/shared/test_nav_scope_filter.py similarity index 95% rename from packages/shared-python/shared/tests/test_nav_scope_filter.py rename to deprecated/mapnav/tests/shared/test_nav_scope_filter.py index 2228d362..ecf02fc9 100644 --- a/packages/shared-python/shared/tests/test_nav_scope_filter.py +++ b/deprecated/mapnav/tests/shared/test_nav_scope_filter.py @@ -5,13 +5,14 @@ import json from typing import Any, List -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( +from shared.services.retrieval.nav.nav_knowhere import NamespaceKnowhereProvider +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, - NamespaceKnowhereProvider, SectionRow, ) -from shared.services.retrieval.nav.nav_node_filter import field_predicate, node_filter +from shared.services.retrieval.nav.nav_node_filter import node_filter +from shared.services.retrieval.scoring.node_filter_predicates import field_predicate from shared.services.retrieval.nav.nav_scope_filter import run_scope_filter from shared.services.retrieval.nav.nav_types import NavConfig diff --git a/packages/shared-python/shared/tests/test_nav_snapshot.py b/deprecated/mapnav/tests/shared/test_nav_snapshot.py similarity index 97% rename from packages/shared-python/shared/tests/test_nav_snapshot.py rename to deprecated/mapnav/tests/shared/test_nav_snapshot.py index f8bf76f6..76911ac1 100644 --- a/packages/shared-python/shared/tests/test_nav_snapshot.py +++ b/deprecated/mapnav/tests/shared/test_nav_snapshot.py @@ -13,7 +13,7 @@ import pytest -from shared.services.retrieval.nav.nav_knowhere import SectionRow, UnitRow +from shared.services.retrieval.scoring.knowhere_provider import SectionRow, UnitRow from shared.services.retrieval.nav_snapshot import build_nav_snapshot diff --git a/packages/shared-python/shared/tests/test_nav_stamp.py b/deprecated/mapnav/tests/shared/test_nav_stamp.py similarity index 100% rename from packages/shared-python/shared/tests/test_nav_stamp.py rename to deprecated/mapnav/tests/shared/test_nav_stamp.py diff --git a/packages/shared-python/shared/tests/test_nav_trace_map.py b/deprecated/mapnav/tests/shared/test_nav_trace_map.py similarity index 100% rename from packages/shared-python/shared/tests/test_nav_trace_map.py rename to deprecated/mapnav/tests/shared/test_nav_trace_map.py diff --git a/packages/shared-python/shared/services/retrieval/trace/mapnav.py b/deprecated/mapnav/trace_mapnav.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/trace/mapnav.py rename to deprecated/mapnav/trace_mapnav.py diff --git a/docs/assets/knowhere-banner-2.0.png b/docs/assets/knowhere-banner-2.0.png new file mode 100644 index 00000000..6332a6ee Binary files /dev/null and b/docs/assets/knowhere-banner-2.0.png differ diff --git a/docs/assets/step-2-agent-neutral-retrieval.png b/docs/assets/step-2-agent-neutral-retrieval.png new file mode 100644 index 00000000..21e075d6 Binary files /dev/null and b/docs/assets/step-2-agent-neutral-retrieval.png differ diff --git a/docs/design/entity-node-graph.md b/docs/design/entity-node-graph.md new file mode 100644 index 00000000..3fe6540a --- /dev/null +++ b/docs/design/entity-node-graph.md @@ -0,0 +1,133 @@ +# Chunk-level entity graph (deferred) + +**Status:** Deferred — not scheduled, no code written. This replaces the +old "entity 节点 / 共现边：schema 文档预留，发布侧不改" non-goal line in +`.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md` §四 with +an actual future phase (see that plan's todo list for the tracking id). + +**Relation to Phase 2 (agent_tools registry)**: does not block or change +it. Reasoning below in "Why this doesn't affect Phase 2." + +## Problem + +Today `DocumentGraphService.publish_document_graph` +(`packages/shared-python/shared/services/retrieval/graph/service.py:70`) +only ever creates **one graph node per document** and only ever creates +**document ↔ document** `related` edges. The entity/keyword overlap that +drives those edges is computed by first collapsing *every chunk in the +document* into one flat entity set +(`get_normalized_entity_set(chunk_metadata_list)`, `service.py:100`) before +any comparison happens — so by the time two documents get connected, you +have already lost which section/chunk in document A actually shares an +entity with which section/chunk in document B. + +The ask: connect the actual **chunks** that share an entity, not just the +two documents they happen to live in. + +## What "node" means here (resolved) + +A chunk is already, in the overwhelming majority of cases, the +bottom-level unit — one leaf section, one body chunk. The one documented +exception (see `CORPUS_SCHEMA.md` §2, and verified against real published +data: 170 of 801 sections in a two-document local corpus) is **`chunk` +track (text-track) only**: a section that has its own content directly +under its heading, before its first child heading, owns a body chunk of +its own even though it also has children. `page_memory` (page) track has +no such exception — only leaves ever own a body chunk there. + +So there is no separate "section-level" option distinct from "chunk-level" +— they are the same thing, because a section owns *at most one* body +chunk. The node for this graph is the chunk: every `text`/`page` body +chunk, and — no reason to special-case them out — every `image`/`table` +chunk too, since they already carry their own `entities` in +`chunk_metadata` independent of which section they are `connect_to`-linked +from. Structural sections that own no chunk of their own simply have no +entities and never become nodes; no rollup needed, no ambiguity. + +**Naming collision to keep in mind when this is built**: the Phase 2 +`node_filter` tool's "node" means a node in the *section/hierarchy tree* +(operates on `section_path` and `summary`, see `CORPUS_SCHEMA.md` §6). The +graph's "node" (`graph_nodes` table) is a completely different structure — +today only `node_kind='document'` rows, this proposal adds +`node_kind='chunk'` rows. Do not conflate the two when writing tool +descriptions or code comments for whichever tool eventually exposes this +graph. + +## Why the DB schema needs no migration + +Verified against the actual model +(`packages/shared-python/shared/models/database/document.py:495-575`): + +- `GraphNode.node_kind` is a plain `String(32)`, not an enum constrained to + `'document'` — a new `'chunk'` value needs no schema change. +- `GraphNode.ref_section_id` already exists as a nullable column + (`document.py:515`) and is already set to `None` for document nodes + (`service.py:137`) — it is unused, not absent. +- `GraphEdge.source_node_id` / `target_node_id` are plain FKs to + `graph_nodes.node_id` (`document.py:546-555`) with no constraint that + both ends share a `node_kind` — a chunk-node ↔ chunk-node edge is already + legal today, mechanically. + +So this is additive: keep the existing `node_kind='document'` nodes/edges +exactly as they are (`neighbors` in Phase 2 keeps querying those, unaffected +— see below), and add `node_kind='chunk'` rows and edges alongside them. + +## What actually needs new design (not just "add a node_kind") + +1. **Stop collapsing to one set per document.** Index each qualifying + chunk's own `entities` (`chunk_metadata.entities`, already extracted per + chunk via `extract_entities_from_chunk_metadata` in + `graph/keywords.py:53`) as its own node, instead of merging via + `get_normalized_entity_set` before any comparison. + +2. **Replace the O(other documents) peer loop with an inverted index.** + `service.py:152-161` currently loads every other `node_kind='document'` + row in the namespace and compares against it — fine at "tens/hundreds + of documents" scale. At chunk scale (tens of thousands of chunks per + namespace) this must become an `entity_key → chunk_node_id` lookup + (a dedicated join table with a plain index beats a JSONB containment + scan at this volume) so a newly published chunk only compares against + chunks that already share at least one entity key, not every chunk in + the namespace. + +3. **Retune the overlap threshold — it does not transfer.** + `graph/keywords.py:7-11`: `MIN_ENTITY_OVERLAP = 2`, + `MIN_SCORE_THRESHOLD = 0.8`, and `compute_entity_score` + (`keywords.py:80-96`) is `weight * shared_weight / min(weight_a, + weight_b)`. This was tuned for whole-document aggregate sets (tens of + entities). A single chunk typically carries 1-3 entities (verified: + one real image chunk had exactly 2). Requiring `≥2` shared entities out + of a 1-3 entity set will rarely fire; and when a tiny set *does* overlap + by even one entity, the length-weighted score can trivially hit 1.0. + Neither behavior is useful. Proposed direction: connect on **≥1** shared + entity, but gate on that entity's **rarity across the namespace** + (inverse document frequency over chunks, not documents) so a common + entity (a year, a generic org name) does not wire every chunk to every + other chunk. This namespace-wide chunk-frequency count does not exist + today (`compute_tfidf_keywords` in `keywords.py:100` only computes + document-frequency *within one document's own chunks*, for its + `top_keywords`, not across the namespace) — it is new infrastructure, + not a reuse of an existing utility. + +4. **Only materialize nodes for chunks that have ≥1 entity.** Chunks with + an empty `entities` list should not get a `graph_nodes` row at all, or + the table balloons with rows that can never have an edge. + +5. **Recommend keeping this additive, not a replacement.** Leave + `node_kind='document'` publication in `publish_document_graph` exactly + as-is; add chunk-level publication as a new, separate write path (can + live in the same service or a sibling one). Lower risk, and the + existing document-level `related` edges stay useful for + `list_documents`/`neighbors`-style "what else is like this document" + overviews that don't need chunk precision. + +## Why this doesn't affect Phase 2 + +Phase 2's tool list (`agent_tools/registry.py` and the 8 tools in +`CORPUS_SCHEMA.md` §6) touches the graph only through `neighbors`, which is +scoped to `node_kind='document'` `related` edges — exactly what exists +today, unchanged by this proposal. `node_filter`, `outline`, `recall`, +`grep`, `read`, `assets`, `list_documents` never touch `graph_nodes` / +`graph_edges` at all. This entity-node-graph work is a later, additive +phase with its own exposure decision (new tool? extend `neighbors` with a +granularity param? not decided — out of scope until this is scheduled). diff --git a/docs/design/retrieval-heart-vessel-trace-optimization-plan.md b/docs/design/retrieval-heart-vessel-trace-optimization-plan.md index ed62a57b..f4521260 100644 --- a/docs/design/retrieval-heart-vessel-trace-optimization-plan.md +++ b/docs/design/retrieval-heart-vessel-trace-optimization-plan.md @@ -17,7 +17,7 @@ request, not a historical request: - Trace: `01a05e530f7787333c0e32ca16633b85` - Route: `/api/v1/retrieval/query` - HTTP status: `200` -- Router: `mapnav` +- Router: `agent_explore` - Stop reason: `completed` - Server span: `42.365 s` - Client wall time: `44.690 s` @@ -302,13 +302,13 @@ Acceptance criteria: demonstrated; otherwise this remains a validation note, not an optimization slice; - selected chunk IDs, scores, ordering, and evidence remain unchanged; -- benchmark results record the route family (`classic`, `mapnav`, or +- benchmark results record the route family (`classic`, `agent_explore`, or `small_corpus`) and separate cold, warm, and response-cache-hit requests; cache-hit timings are not mixed into cold-request latency claims. ### P0: Make map-unit projection token-selective -The current map-nav reader already makes its frequency lookup token-selective, +The current map-unit reader already makes its frequency lookup token-selective, but it still loads every revision-scoped map unit before applying the query tokens. Change only the unit projection: start from `document_map_unit_tokens` filtered by `channel` and `token_hash`, then join the diff --git a/docs/design/retrieval-serving-index-rollout-runbook.md b/docs/design/retrieval-serving-index-rollout-runbook.md index 56d2c3cc..2d91c97b 100644 --- a/docs/design/retrieval-serving-index-rollout-runbook.md +++ b/docs/design/retrieval-serving-index-rollout-runbook.md @@ -153,15 +153,16 @@ index with `DROP INDEX CONCURRENTLY`, and rerun the migration. Deploy the application after the additive migrations finish. New publications will write coherent format-v2 statistics. Existing revisions with NULL channel statistics remain on the full scope-first map-unit reader until maintenance -completes; missing, legacy, or unusable indexes remain on the legacy reader. +completes; missing, legacy, or unusable indexes raise. Immediately verify: - API health checks pass; - no migration or model-loading error appears in API logs; -- classic and map-nav requests still complete; +- classic and agent_explore requests still complete; - incomplete-index warnings distinguish statistics-incomplete map-unit serving - from `fallback=legacy_fts`; neither case may return partial or empty results; + from an unusable index, which must raise; statistics-incomplete serving + must not return partial or empty results; - no increase appears in retrieval errors or timeouts. ## Phase 4: Backfill existing format-v2 indexes @@ -317,12 +318,12 @@ Exercise at least: 1. v1 `use_agentic=false`; 2. v2 `use_agentic=false` with equivalent retrieval fields; -3. one `use_agentic=true` map-nav smoke; -4. one request with `use_agentic` omitted, confirming it routes to map-nav; +3. one `use_agentic=true` agent_explore smoke; +4. one request with `use_agentic` omitted, confirming it routes to agent_explore; 5. one filtered request, confirming filtered-scope semantics and the safe fallback where required. -Map-nav LLM output is nondeterministic. For production smoke, require successful +Agent-explore LLM output is nondeterministic. For production smoke, require successful completion, valid citations, expected namespace isolation, and relevant evidence. Do not require byte-identical ordering between independent Planner runs. Deterministic map-score parity remains covered by the contract suite. @@ -339,13 +340,13 @@ recorded baseline, and retrieval error/timeout rates must not regress. - retrieval request p50/p95 and maximum latency, separated by `router_used`; - classic `search.map_unit_discovery` stages: units, frequencies, indexes, statistics, scoring, and hydration; -- map-nav snapshot, episode, and hydration stages; +- agent_explore episode, tool, and hydration stages; - PostgreSQL statement timeouts, lock waits, CPU, I/O, and connection usage; - Redis errors and namespace snapshot cache misses; -- retrieval errors, incomplete-index fallbacks, and response timeouts; -- process CPU and maximum RSS from the corrected map-nav resource log. +- retrieval errors, unusable-index failures, and response timeouts; +- process CPU and maximum RSS from agent_explore resource logs. -Do not mix classic and map-nav latency distributions. Do not treat Redis-warm +Do not mix classic and agent_explore latency distributions. Do not treat Redis-warm snapshot measurements as cold-request performance. ## Pause and resume @@ -365,7 +366,7 @@ If application errors, timeouts, or quality regressions occur: 1. stop the backfill process; 2. redeploy the previous application version; -3. verify classic and map-nav requests using the frozen quality set; +3. verify classic and agent_explore requests using the frozen quality set; 4. retain the additive columns, index, and already-computed statistics unless database health specifically requires their removal. @@ -387,7 +388,7 @@ Attach the following to the deployment ticket: - final `--check` output; - ready/current revision counts; - frozen-query parity results; -- classic and map-nav latency summaries; +- classic and agent_explore latency summaries; - observed fallback, error, and timeout counts; - rollback decision or explicit confirmation that rollback was not required. diff --git a/docs/design/retrieval-streaming-sse.md b/docs/design/retrieval-streaming-sse.md index b0b5a73e..47420fde 100644 --- a/docs/design/retrieval-streaming-sse.md +++ b/docs/design/retrieval-streaming-sse.md @@ -7,7 +7,7 @@ ## Purpose Online Brain users currently wait for a complete retrieval response while -map-nav planning, searching, source review, and final hydration run. This +agentic retrieval, tool calls, and final hydration run. This design makes that work visible without exposing chain-of-thought or changing who owns answer generation. @@ -61,7 +61,7 @@ response as the authoritative result: "sequence": 7, "elapsed_ms": 2410, "status": "completed", - "response": { "namespace": "default", "query": "...", "router_used": "mapnav", "evidence_text": "...", "referenced_chunks": [], "results": [] } + "response": { "namespace": "default", "query": "...", "router_used": "agent_explore", "evidence_text": "...", "referenced_chunks": [], "results": [] } } ``` @@ -90,9 +90,9 @@ terminal SSE event because the HTTP status can no longer be changed. ## Internal implementation seam -Keep the synchronous map-nav implementation. Add an optional callback that -receives a sanitized progress projection after each completed planner, -search, or review step. The SSE route bridges this callback to an +Keep the live retrieval implementation. Add an optional callback that +receives a sanitized progress projection after each completed agent +or classic step. The SSE route bridges this callback to an `asyncio.Queue` using a thread-safe loop handoff while retrieval continues in its existing worker thread. @@ -106,12 +106,12 @@ Add cooperative cancellation checks between steps. An in-flight synchronous provider call may finish before cancellation takes effect. Cancelled runs do not perform final hydration when cancellation is observed in time. -Phase ownership is explicit: the route emits `started`; the map-nav adapter -emits `planning` before `plan_query` and `searching` before navigation or -classic discovery; the route emits `reviewing_sources` after retrieval -selection and before reference hydration; and it emits `finalizing` before -public projection. Counts are sourced from existing snapshot, reference, and -assembled-result counts and are omitted when not yet known. +Phase ownership is explicit: the route emits `started`; the live route +emits `searching` during classic discovery or agent_explore tool steps; the +route emits `reviewing_sources` after retrieval selection and before +reference hydration; and it emits `finalizing` before public projection. +Counts are sourced from existing reference and assembled-result counts and +are omitted when not yet known. ## Correct duration accounting @@ -132,15 +132,15 @@ Required changes: - pass the execution start timestamp into `TraceRecorder`; - set `retrieval_runs.latency_ms` from that timestamp; - record cache-hit runs with the same definition; -- ensure classic, map-nav, small-corpus, cache-hit, failed, and cancelled +- ensure classic, agent_explore, small-corpus, cache-hit, failed, and cancelled retrievals all have an explicit timing/observability outcome; - expose separate `time_to_first_event_ms`, `retrieval_latency_ms`, and downstream `time_to_first_token_ms` measurements; - retain per-step `elapsed_ms` as step latency, not total request latency. -`retrieval_runs` is the ledger for every retrieval execution, not only -map-nav. Each row records the route type, `agentic_enabled`, `cache_hit`, -canonical latency, and terminal status for classic, map-nav, small-corpus, +`retrieval_runs` is the ledger for every retrieval execution. Each row +records the route type, `agentic_enabled`, `cache_hit`, +canonical latency, and terminal status for classic, agent_explore, small-corpus, cache-hit, failed, and cancelled runs. Add a backward-compatible status field and migration rather than overloading free-form error text. @@ -186,7 +186,7 @@ later local constructor time or include its own flush duration. ## Verification gates -- correct phase order for map-nav, classic, and small-corpus routes; +- correct phase order for agent_explore, classic, and small-corpus routes; - cache-hit streams emit only applicable phases and identify the cache hit; - no sensitive planner or evidence data before the terminal event; - terminal citations match the existing JSON endpoint; diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 6a46fdd2..5b847ab9 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -7,7 +7,6 @@ from uuid import uuid4 from sqlalchemy import ( - Computed, JSON, DateTime, Float, @@ -20,7 +19,6 @@ Text, UniqueConstraint, ) -from sqlalchemy.dialects.postgresql import TSVECTOR from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base @@ -176,22 +174,6 @@ class DocumentChunk(Base): content_search_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) path_search_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) term_search_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - content_search_tsv: Mapped[Optional[str]] = mapped_column( - TSVECTOR, - Computed( - "to_tsvector('simple', COALESCE(content_search_text, ''))", - persisted=True, - ), - nullable=True, - ) - path_search_tsv: Mapped[Optional[str]] = mapped_column( - TSVECTOR, - Computed( - "to_tsvector('simple', COALESCE(path_search_text, ''))", - persisted=True, - ), - nullable=True, - ) source_chunk_path: Mapped[Optional[str]] = mapped_column(Text, nullable=True) file_path: Mapped[Optional[str]] = mapped_column(Text, nullable=True) chunk_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column( @@ -232,16 +214,6 @@ class DocumentChunk(Base): "id", ), Index("idx_document_chunks_section", "section_id"), - Index( - "idx_chunk_content_search_tsv", - "content_search_tsv", - postgresql_using="gin", - ), - Index( - "idx_chunk_path_search_tsv", - "path_search_tsv", - postgresql_using="gin", - ), ) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py b/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py new file mode 100644 index 00000000..d313f765 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/__init__.py @@ -0,0 +1,23 @@ +"""In-process tool-calling agentic retrieval route. + +Default agentic path when ``use_agentic`` is unset or true. Runs the +``agent_tools`` corpus registry through an LLM tool-calling loop. Map-nav is +archived under ``deprecated/mapnav/`` and is not a live route. + +Which provider runs that loop (Cursor SDK or OpenAI-compatible) is the +``AGENT_EXPLORE_HARNESS`` switch resolved via ``resolve_harness()``. +""" + +from __future__ import annotations + +from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_explore.harness import Harness, resolve_harness +from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult + +__all__ = [ + "AgentStep", + "EpisodeBudget", + "EpisodeResult", + "Harness", + "resolve_harness", +] diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py b/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py new file mode 100644 index 00000000..4f840426 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py @@ -0,0 +1,47 @@ +"""Bridge from an ``agent_explore`` episode to the existing decision-trace shape. + +``DecisionTraceStep`` / ``TraceRecorder`` (``shared/services/retrieval/trace/``) +are already provider-agnostic — this module only maps this package's own +``AgentStep`` records onto that shared shape. +""" + +from __future__ import annotations + +from shared.services.retrieval.agent_explore.types import AgentStep +from shared.services.retrieval.trace import DecisionTraceStep + +# Cap raw trace text before it goes into the public decision_trace response. +TRACE_OBSERVATION_MAX_CHARS = 2_000 + + +def build_decision_trace(steps: list[AgentStep]) -> list[DecisionTraceStep]: + trace_steps: list[DecisionTraceStep] = [] + for step in steps: + observation_text = step.observation_text + if len(observation_text) > TRACE_OBSERVATION_MAX_CHARS: + observation_text = observation_text[:TRACE_OBSERVATION_MAX_CHARS] + "..." + phase = "finish" if step.tool_name == "finish" else ( + "stop" if not step.tool_name else "tool_call" + ) + trace_steps.append( + DecisionTraceStep( + step_index=step.step_index, + agent="agent_explore", + phase=phase, + observation={"observation_text": observation_text}, + decision={ + "action": step.tool_name or "no_tool_call", + "args": step.tool_args, + }, + result={ + "status": "error" if step.error else "ok", + "error": step.error, + }, + budget={ + "tokens_used_delta": step.tokens_used_delta, + "tokens_used_total": step.tokens_used_total, + }, + elapsed_ms=step.elapsed_ms, + ) + ) + return trace_steps diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/budget.py b/packages/shared-python/shared/services/retrieval/agent_explore/budget.py new file mode 100644 index 00000000..47a0c606 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/budget.py @@ -0,0 +1,80 @@ +"""Per-episode budget enforcement for ``agent_explore``. + +Standalone from ``nav/nav_token_budget.py`` on purpose — this package must +have no runtime dependency on ``nav/`` (see ``config.py``'s module +docstring). Re-reads the same ``RETRIEVAL_NAV_TOKEN_LIMIT`` env var the plan +calls for, so operators keep one token-limit knob across both agentic +routes, but the counting mechanism is a fresh, request-scoped object (this +episode runs as one async function, not a separate thread with recursive +calls), not the nav contextvar machinery ``nav_token_budget`` needed for its +own call shape. +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass, field +from typing import Any + +from shared.services.retrieval.agent_explore.config import ( + AGENT_EXPLORE_MAX_STEPS, + AGENT_EXPLORE_WALL_CLOCK_SECONDS, +) + +_ENV_TOKEN_LIMIT = "RETRIEVAL_NAV_TOKEN_LIMIT" +_DEFAULT_TOKEN_LIMIT = 100_000 + +StopReason = str # one of: "token_limit" | "max_steps" | "wall_clock" + + +def resolve_token_limit() -> int: + """Always a positive limit: env override, else the shared default.""" + try: + limit = int(os.environ.get(_ENV_TOKEN_LIMIT, "").strip()) + except ValueError: + limit = 0 + return limit if limit > 0 else _DEFAULT_TOKEN_LIMIT + + +@dataclass +class EpisodeBudget: + """Tracks one episode's LLM-token / step / wall-clock spend.""" + + token_limit: int = field(default_factory=resolve_token_limit) + max_steps: int = AGENT_EXPLORE_MAX_STEPS + wall_clock_seconds: float = AGENT_EXPLORE_WALL_CLOCK_SECONDS + tokens_used: int = 0 + steps_used: int = 0 + _started_at: float = field(default_factory=time.monotonic, repr=False) + + def record_usage(self, usage: dict[str, Any] | None) -> None: + try: + add = int((usage or {}).get("total_tokens", 0) or 0) + except (TypeError, ValueError): + add = 0 + if add > 0: + self.tokens_used += add + + def record_step(self) -> None: + self.steps_used += 1 + + def exhausted(self) -> StopReason | None: + """Return which budget dimension is exceeded, if any, else None.""" + if self.tokens_used >= self.token_limit: + return "token_limit" + if self.steps_used >= self.max_steps: + return "max_steps" + if time.monotonic() - self._started_at >= self.wall_clock_seconds: + return "wall_clock" + return None + + def snapshot(self) -> dict[str, Any]: + return { + "token_limit": self.token_limit, + "tokens_used": self.tokens_used, + "max_steps": self.max_steps, + "steps_used": self.steps_used, + "wall_clock_seconds": self.wall_clock_seconds, + "elapsed_seconds": round(time.monotonic() - self._started_at, 3), + } diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/config.py b/packages/shared-python/shared/services/retrieval/agent_explore/config.py new file mode 100644 index 00000000..0d467633 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/config.py @@ -0,0 +1,110 @@ +"""Production config for the ``agent_explore`` in-process tool-loop. + +Model choice for the OpenAI-compatible harness is ``deepseek-v4-flash``: +tool-calling (single, parallel, and forced ``tool_choice``) was verified +live against this exact model; no other model has been verified for this +codebase's OpenAI-compatible client. +""" + +from __future__ import annotations + +AGENT_EXPLORE_MODEL = "deepseek-v4-flash" + +# Model for AGENT_EXPLORE_HARNESS=cursor_sdk (harness/cursor_harness.py) — +# a separate constant from AGENT_EXPLORE_MODEL because that harness's +# provider (Cursor SDK) is a disjoint model catalog from the OpenAI-compatible +# client's, not an interchangeable choice. "composer-2.5" is what the PoC +# (apps/worker/scripts/debug_cursor_agent_explore.py) verified live and +# noticeably outperformed deepseek-v4-flash on the two hardest eval-fixture +# queries (q04/q06) — see the Phase 3.5 landing record. +AGENT_EXPLORE_CURSOR_MODEL = "composer-2.5" + +# One LLM turn = one round-trip that may contain several parallel tool calls +# (see episode.py). +AGENT_EXPLORE_MAX_STEPS = 12 + +# Wall-clock ceiling for the whole episode (LLM round-trips + tool +# dispatch), independent of the token budget. New constant, not tuned yet. +AGENT_EXPLORE_WALL_CLOCK_SECONDS = 180.0 + +# Max tokens requested per LLM completion turn (thinking disabled — see +# episode.py; this is completion budget, not context window). +AGENT_EXPLORE_MAX_COMPLETION_TOKENS = 1024 + +FINISH_TOOL_NAME = "finish" + +FINISH_TOOL_SCHEMA: dict[str, object] = { + "type": "object", + "properties": { + "refs": { + "type": "array", + "description": ( + "Final cited evidence, in priority order. Each item " + "identifies one section or chunk you have already looked at " + "via corpus.read (or, for an asset, corpus.assets)." + ), + "items": { + "type": "object", + "properties": { + "document_id": {"type": "string"}, + "section_path": {"type": "string"}, + "chunk_id": {"type": "string"}, + }, + "required": ["document_id"], + }, + }, + "notes": { + "type": "string", + "description": ( + "Optional short note on why these refs answer the query, " + "or why none were found." + ), + }, + }, + "required": ["refs"], +} + +FINISH_TOOL_DESCRIPTION = ( + "Call this when you have gathered enough evidence to answer the query, " + "or when you are certain the corpus does not contain an answer. This " + "ends the exploration — do not answer in plain text; the final answer " + "is synthesized downstream from the refs you cite here." +) + +# Appended after the verbatim CORPUS_SCHEMA.md text (schema_doc.py) to form +# this harness's system prompt. Kept out of CORPUS_SCHEMA.md itself because +# the finish-tool loop contract is agent_explore-specific, not something the +# MCP-facing harnesses (Cursor/Codex/Claude) need — see that file's own +# "single source... do not duplicate" header. +LOOP_CONTRACT_SUFFIX = f""" + +--- + +## Exploration loop contract + +You are exploring this corpus autonomously to answer one query. Use the +tools above to navigate; you may call several tools in one turn when they +are independent. When you have enough evidence, call `{FINISH_TOOL_NAME}` +with the `refs` you want cited as the answer — do not write the final answer +as plain text yourself, it is synthesized downstream from your cited refs. +If you exhaust your tool budget without a confident answer, call +`{FINISH_TOOL_NAME}` with your best-effort `refs` (or an empty list plus a +`notes` explanation of why nothing was found) rather than continuing to +call other tools. + +`refs` is REQUIRED and must not be omitted or left empty if you called +`corpus.read` (or `corpus.assets`) even once during this exploration: copy +the `document_id` and `chunk_id`/`section_path` of every section/chunk you +read that supports your answer into `refs` before calling +`{FINISH_TOOL_NAME}`. Calling `{FINISH_TOOL_NAME}` with no `refs` after +having already read relevant content discards that evidence. + +Before calling `{FINISH_TOOL_NAME}`, if you have not called `corpus.read` +(or `corpus.assets`) even once this exploration, you have not actually +verified anything yet — a search tool returning candidates is not the same +as having read them. In that case, either read your best candidate section +first, or — only if you have positively confirmed there is nothing to read +(e.g. a structural check came back with zero matching sections) — say so +explicitly in `notes`. Repeatedly rephrasing the same search instead of +reading a candidate you already found is not a substitute for reading it. +""" diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py new file mode 100644 index 00000000..22621db5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py @@ -0,0 +1,56 @@ +"""Concurrency-safe per-call tool dispatch shared by every ``agent_explore`` harness. + +Mirrors ``apps/api/app/mcp/dynamic_tools.py``'s ``_dispatch_tool``: SQLAlchemy's +``AsyncSession`` is not safe for concurrent use from multiple coroutines, and a +harness cannot always control whether the orchestrating model issues tool +calls in parallel (verified against the Cursor SDK harness — see +``harness/cursor_harness.py``'s module docstring). Opening a short-lived +session per tool call, instead of sharing one ``AsyncSession`` across the +whole episode, makes dispatch safe regardless of how a harness calls it — +strictly sequential (``harness/openai_harness.py``) or genuinely concurrent +(``harness/cursor_harness.py``). + +This is a behavior change for the OpenAI harness too (previously one shared +session for the whole episode), but a safe one: it already dispatched +sequentially, so a fresh session per call only adds one extra connection +checkout per tool call, never a correctness risk. +""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.agent_tools import REGISTRY, ToolBudget, ToolContext, ToolResult + +DbFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] + + +async def dispatch_tool_call( + name: str, + args: dict[str, Any], + *, + db_factory: DbFactory, + user_id: str, + namespace: str, + budget: ToolBudget | None = None, +) -> ToolResult: + """Run one ``REGISTRY`` tool call against a fresh, call-scoped DB session. + + ``budget`` is forwarded into the ``ToolContext`` built for this one call + (defaults to ``ToolBudget()``, matching prior behavior); callers that need + the same budget value for their own text-capping (``shared.tool_message_content``) + should hold onto the ``ToolBudget`` they pass here rather than reach back + into the (call-scoped, already-closed) ``ToolContext``. + """ + try: + async with db_factory() as db: + tool_ctx = ToolContext( + db=db, user_id=user_id, namespace=namespace, budget=budget or ToolBudget() + ) + return await REGISTRY.dispatch(name, tool_ctx, args) + except Exception as exc: # noqa: BLE001 - one broken tool must not kill the episode + return ToolResult(text="", error=f"{type(exc).__name__}: {exc}") diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/__init__.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/__init__.py new file mode 100644 index 00000000..e4559f30 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/__init__.py @@ -0,0 +1,14 @@ +"""Pluggable ``agent_explore`` harness implementations. + +``base.Harness`` is the provider-agnostic interface; ``openai_harness`` and +``cursor_harness`` are the two implementations selected via +``AGENT_EXPLORE_HARNESS`` (``execution/routes.py``). See the Phase 3.5 +section of ``.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md``. +""" + +from __future__ import annotations + +from shared.services.retrieval.agent_explore.harness.base import Harness +from shared.services.retrieval.agent_explore.harness.resolve import resolve_harness + +__all__ = ["Harness", "resolve_harness"] diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/base.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/base.py new file mode 100644 index 00000000..4bab470e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/base.py @@ -0,0 +1,42 @@ +"""Provider-agnostic ``Harness`` interface for ``agent_explore``. + +Any backend-hosted tool-loop implementation (OpenAI-compatible/DeepSeek, +Cursor SDK, ...) implements this single method. ``_run_agent_explore_route`` +(``execution/routes.py``) resolves one ``Harness`` via ``AGENT_EXPLORE_HARNESS`` +(see ``harness/resolve.py``) and calls it the same way regardless of which +provider is behind it — the route/bridge/budget-object shapes do not change +per harness. +""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_explore.dispatch import DbFactory +from shared.services.retrieval.agent_explore.types import EpisodeResult + + +@runtime_checkable +class Harness(Protocol): + """One backend-hosted tool-loop implementation over ``agent_tools.REGISTRY``.""" + + async def run_episode( + self, + *, + db_factory: DbFactory, + user_id: str, + namespace: str, + query: str, + budget: EpisodeBudget, + ) -> EpisodeResult: + """Explore the corpus for ``query`` and return the cited evidence. + + ``db_factory`` is a call-scoped DB session factory (see + ``dispatch.py``) — implementations must not hold one shared + ``AsyncSession`` across the whole episode; every ``REGISTRY.dispatch`` + call goes through ``dispatch.dispatch_tool_call(..., db_factory=db_factory)`` + so concurrent tool calls (a real Cursor SDK behavior, not just a + theoretical one — see ``harness/cursor_harness.py``) are always safe. + """ + raise NotImplementedError diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py new file mode 100644 index 00000000..6678be43 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py @@ -0,0 +1,302 @@ +"""``Harness`` implementation over the Cursor SDK local agent. + +Promoted from ``apps/worker/scripts/debug_cursor_agent_explore.py`` (PoC) — +see Phase 3.5 of +``.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md``. Uses +``cursor_sdk``'s ``local.custom_tools`` so this host process still executes +``agent_tools.REGISTRY`` (same ``corpus.*`` tools as ``harness/openai_harness.py``), +while the orchestration model comes from Cursor (default +``config.AGENT_EXPLORE_CURSOR_MODEL``, e.g. ``composer-2.5``) instead of +DeepSeek. + +Requires ``cursor-sdk`` (base dependency of ``apps/api``; worker debug +scripts use the ``cursor-harness`` extra) and ``CURSOR_API_KEY``. The +guarded import in ``_require_cursor_sdk`` raises a clear error at episode +start if ``cursor-sdk`` isn't installed. + +Architectural difference from ``openai_harness.py`` that budget enforcement +has to work around: this harness does not control the LLM turn loop. +``agent.send(...)`` + ``await run.wait()`` hands the *entire* multi-turn +tool-calling loop to the Cursor SDK; the host process only sees (a) +``execute`` callbacks for each ``corpus.*``/``finish`` tool call — run +synchronously off the SDK's own thread and bridged back onto this event +loop via ``asyncio.run_coroutine_threadsafe`` (mirrors the PoC's +``_dispatch_sync``) — and (b) the terminal ``RunResult`` once ``wait()`` +returns. There is no per-turn hook to inspect budget mid-turn and force +``tool_choice`` the way ``openai_harness.py`` does. Each budget dimension is +therefore enforced (or explicitly not) differently here: + +- **``max_steps``**: a plain counter (``budget.steps_used``, incremented once + per ``corpus.*`` dispatch — ``finish`` does not count, it ends the episode + on its own) checked in ``_dispatch_sync`` *before* dispatching. Once the + counter reaches ``budget.max_steps``, further ``corpus.*`` calls are not + forwarded to ``REGISTRY`` at all — the callback returns a fixed + "budget exhausted, call finish now" string instead. This is a real, + synchronous cutoff (unlike the two dimensions below): no reliance on + cancelling the SDK run from a background task. +- **``wall_clock``**: ``asyncio.wait_for(run.wait(), timeout=budget.wall_clock_seconds)``, + per the plan. On timeout, best-effort ``await run.cancel()`` (own short + timeout, so a hung cancel RPC can't hang this call forever) so the + underlying agent run actually stops server-side instead of merely being + abandoned by this process, then the episode result is built from whatever + ``finish``/tool-trajectory refs were captured before the timeout. +- **``token_limit``**: **not actively enforced mid-run.** Verified by reading + the installed ``cursor-sdk`` package's source + (``cursor_sdk._run_base._RunBase.usage``, 2026-09) that cumulative token + usage is incrementally accumulated from ``SDKUsageMessage`` stream events + as ``run.wait()`` consumes them, and exposed as a live ``run.usage`` + property — so a mid-run cutoff (e.g. a polling task calling + ``run.cancel()``) is *plausible*. It was **not** exercised against a real + ``CURSOR_API_KEY`` run in this change (no key available in the + implementation environment), so building an active cutoff on top of an + unverified live-update assumption was judged higher-risk than shipping. + ``EpisodeBudget.exhausted()``'s ``token_limit`` dimension is therefore left + unchecked here by design; ``budget.tokens_used`` is only populated + *post-hoc* from the terminal ``RunResult.usage`` for observability + (``EpisodeResult.tokens_used``), after the episode has already finished. + If ``eval-cursor-harness`` confirms ``run.usage`` does update live against + a real key, promoting this to an active cutoff (mirroring the + ``max_steps``/``wall_clock`` pattern above) is straightforward follow-up + work — deliberately not done speculatively here. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import time +from typing import Any + +from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_explore.config import ( + AGENT_EXPLORE_CURSOR_MODEL, + FINISH_TOOL_DESCRIPTION, + FINISH_TOOL_NAME, + FINISH_TOOL_SCHEMA, + LOOP_CONTRACT_SUFFIX, +) +from shared.services.retrieval.agent_explore.dispatch import DbFactory, dispatch_tool_call +from shared.services.retrieval.agent_explore.shared import ( + EVIDENCE_TOOL_NAMES, + budget_status_line, + dedup_refs, + normalize_finish_refs, + tool_message_content, + wire_safe_tool_name, +) +from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult +from shared.services.retrieval.agent_tools import ( + REGISTRY, + ToolBudget, + ToolResult, + load_corpus_schema_text, +) + +# Grace period for the underlying agent run to actually stop, after a +# best-effort run.cancel() following a wall_clock timeout, before this +# process gives up waiting for a terminal RunResult and falls back to +# whatever refs the trajectory already captured. Not tuned against real +# cancel-RPC latency yet (no CURSOR_API_KEY in the implementation +# environment) — revisit with eval-cursor-harness data. +_CANCEL_GRACE_SECONDS = 30.0 + +_BUDGET_EXHAUSTED_MESSAGE = ( + "error: step budget exhausted for this episode — do not call any more " + "corpus.* tools; call finish now with whatever refs you already have " + "(or an empty list plus a notes explanation)." +) + + +def _require_cursor_sdk() -> Any: + try: + import cursor_sdk + except ImportError as exc: + raise RuntimeError( + "AGENT_EXPLORE_HARNESS=cursor_sdk requires the " + "'cursor-sdk' dependency, which is not installed in this " + "interpreter. For the API service it is a base dependency " + "(apps/api). Worker debug scripts: uv sync --extra cursor-harness." + ) from exc + return cursor_sdk + + +class CursorHarness: + """``Harness`` implementation over the Cursor SDK local agent.""" + + def __init__(self, *, model: str | None = None) -> None: + self._model = model or AGENT_EXPLORE_CURSOR_MODEL + + async def run_episode( + self, + *, + db_factory: DbFactory, + user_id: str, + namespace: str, + query: str, + budget: EpisodeBudget, + ) -> EpisodeResult: + cursor_sdk = _require_cursor_sdk() + + api_key = os.environ.get("CURSOR_API_KEY", "").strip() + if not api_key: + raise RuntimeError( + "AGENT_EXPLORE_HARNESS=cursor_sdk requires CURSOR_API_KEY to be set" + ) + + tool_budget = ToolBudget() + loop = asyncio.get_running_loop() + + steps: list[AgentStep] = [] + trajectory_refs: list[dict[str, Any]] = [] + # None until finish is actually called — distinguishes "finish + # called with an empty refs list" (respect it) from "finish never + # called" (fall back to trajectory_refs below), same contract as + # openai_harness.py's normalize_finish_refs + fallback. + finish_state: dict[str, Any] = {"refs": None, "notes": ""} + stop_reason = "finished" + + def _dispatch_sync(tool_name: str, args: dict[str, Any]) -> str: + if budget.steps_used >= budget.max_steps: + steps.append( + AgentStep( + step_index=len(steps), + tool_name=tool_name, + tool_args=args, + observation_text=_BUDGET_EXHAUSTED_MESSAGE, + error="budget_max_steps", + elapsed_ms=0, + tokens_used_delta=0, + tokens_used_total=budget.tokens_used, + ) + ) + return _BUDGET_EXHAUSTED_MESSAGE + budget.record_step() + tool_started = time.perf_counter() + future = asyncio.run_coroutine_threadsafe( + dispatch_tool_call( + tool_name, + args, + db_factory=db_factory, + user_id=user_id, + namespace=namespace, + budget=tool_budget, + ), + loop, + ) + try: + tool_result = future.result(timeout=180) + except Exception as exc: # noqa: BLE001 - one broken tool must not kill the episode + tool_result = ToolResult(text="", error=f"{type(exc).__name__}: {exc}") + elapsed_ms = int((time.perf_counter() - tool_started) * 1000) + content = tool_message_content(tool_result, max_chars=tool_budget.max_chars) + # Appended per call, unlike openai_harness.py's once-per-turn + # placement — this harness has no batched-turn concept exposed to + # the host process (see module docstring): each corpus.* dispatch + # is the only per-step hook available to surface budget state. + content_with_budget = content + "\n" + budget_status_line(budget) + steps.append( + AgentStep( + step_index=len(steps), + tool_name=tool_name, + tool_args=args, + observation_text=content_with_budget, + error=tool_result.error, + elapsed_ms=elapsed_ms, + tokens_used_delta=0, + tokens_used_total=budget.tokens_used, + ) + ) + if tool_name in EVIDENCE_TOOL_NAMES and not tool_result.error: + trajectory_refs.extend(tool_result.refs) + return content_with_budget + + custom_tools: dict[str, Any] = {} + for spec in REGISTRY.all(): + wire_name = wire_safe_tool_name(spec.name) + + def _make_execute(resolved_name: str): + def execute(args: dict[str, Any], _ctx: Any) -> str: + return _dispatch_sync(resolved_name, dict(args or {})) + + return execute + + custom_tools[wire_name] = cursor_sdk.CustomTool( + execute=_make_execute(spec.name), + description=f"{spec.description} (canonical name: {spec.name})", + input_schema=spec.json_schema, + ) + + def finish_execute(args: dict[str, Any], _ctx: Any) -> str: + finish_state["refs"] = normalize_finish_refs(args.get("refs")) + finish_state["notes"] = str(args.get("notes") or "") + return json.dumps({"status": "finished", "refs": len(finish_state["refs"])}) + + custom_tools[FINISH_TOOL_NAME] = cursor_sdk.CustomTool( + execute=finish_execute, + description=FINISH_TOOL_DESCRIPTION, + input_schema=FINISH_TOOL_SCHEMA, + ) + + system_prompt = load_corpus_schema_text() + LOOP_CONTRACT_SUFFIX + user_prompt = ( + f"{system_prompt}\n\n---\n\nUser query:\n{query}\n\n" + "Tool names on the wire use underscores " + f"({', '.join(sorted(custom_tools))}). Explore with those tools, " + "then call finish with cited refs." + ) + + result: Any = None + async with await cursor_sdk.AsyncClient.launch_bridge( + workspace=os.getcwd(), + ) as client: + async with await client.agents.create( + cursor_sdk.AgentOptions( + api_key=api_key, + model=self._model, + local=cursor_sdk.LocalAgentOptions( + cwd=os.getcwd(), + custom_tools=custom_tools, + ), + ) + ) as agent: + run = await agent.send(user_prompt) + try: + result = await asyncio.wait_for( + run.wait(), timeout=budget.wall_clock_seconds + ) + except asyncio.TimeoutError: + stop_reason = "budget_wall_clock" + with contextlib.suppress(Exception): + await asyncio.wait_for(run.cancel(), timeout=10) + with contextlib.suppress(Exception): + result = await asyncio.wait_for( + run.wait(), timeout=_CANCEL_GRACE_SECONDS + ) + + result_usage_total_tokens = 0 + if result is not None and result.usage is not None: + result_usage_total_tokens = result.usage.total_tokens + budget.record_usage({"total_tokens": result_usage_total_tokens}) + + result_refs = finish_state["refs"] + result_notes = str(finish_state["notes"] or "") + if not result_refs: + fallback_refs = dedup_refs(trajectory_refs) + if fallback_refs: + result_refs = fallback_refs + result_notes = (result_notes + " " if result_notes else "") + ( + "[refs auto-filled from corpus.read/corpus.assets trajectory; " + "finish did not cite any]" + ) + result_refs = result_refs or [] + + return EpisodeResult( + refs=result_refs, + notes=result_notes, + steps=steps, + stop_reason=stop_reason, + tokens_used=budget.tokens_used, + model_name=self._model, + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py new file mode 100644 index 00000000..80fd5716 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py @@ -0,0 +1,388 @@ +"""In-process tool-calling episode over the ``agent_tools`` corpus registry, +via an OpenAI-compatible (DeepSeek) function-calling loop. + +Uses ``OpenAICompatibleClientSync.chat_completion_raw_with_usage`` — the RAW +response, not ``chat_completion_with_usage`` — because the latter only +returns ``.content`` and silently drops ``message.tool_calls``. Verified live +(2026-09-08) against ``deepseek-v4-flash`` via this codebase's client: +single tool call, parallel tool calls in one turn, tool-result feedback + +final synthesis, and forced ``tool_choice`` (used for the budget-exhaustion +cutoff below) all work. + +Tool calls within one turn are dispatched sequentially through +``dispatch.dispatch_tool_call`` (fresh DB session per call — safe for any +harness, not just this one), not via ``asyncio.gather``: batching several +tool calls into one LLM turn already removes the LLM round-trip per tool +(the dominant cost); true DB-level concurrency within a turn is not +implemented here since this provider's tool calls are handled one at a time +by design, not because concurrent dispatch would be unsafe (it no longer is +— see ``dispatch.py``). + +No import from archived map-nav modules. + +Two Phase 4 fixes (audited live against the eval fixture in +``apps/worker/scripts/fixtures/changheba_archive_eval_queries.json``), both +now shared with any other harness via ``shared.py`` except stale-message +collapsing (fix 1), which stays here — it mutates this harness's own +``messages: list[dict]`` history, a mechanism the Cursor SDK harness has no +equivalent hook for: + +1. **Stale tool-message collapsing** (``_TOOL_MESSAGE_FRESH_TURNS``, + ``_collapse_stale_tool_messages``): ``messages`` only ever appended, so a + single ``corpus.outline``/``corpus.node_filter`` call (each capped at + ``ToolBudget.max_chars`` — currently ``EVIDENCE_TEXT_CHAR_BUDGET=12_000``, + see ``registry.py``) was resent in full on every later turn. Verified + live: two independent queries (q04, q06 in the eval fixture) hit + ``RETRIEVAL_NAV_TOKEN_LIMIT`` (100k default) within 7-8 LLM turns from + this resend alone, not from query difficulty — per-turn token cost grew + monotonically (q04: 4.4k -> 4.8k -> 19.8k -> 21.5k -> 24.2k -> 29.6k). +2. **Trajectory refs fallback** (``shared.dedup_refs`` + the fallback at the + end of ``run_episode``): verified live that ``finish`` can be called with + no ``refs`` key at all (raw ``function.arguments`` was literally ``'{}'``) + even after the model had already read clearly relevant sections via + ``corpus.read`` — the ``FINISH_TOOL_SCHEMA``'s ``"required": ["refs"]`` is + a schema hint, not a provider-enforced constraint. When ``finish``'s own + ``refs`` end up empty (whether from this, from ``no_tool_call``, or from a + forced-finish the provider ignored — all three exit paths), the episode + now falls back to the refs already returned by every + ``corpus.read``/``corpus.assets`` call in the trajectory, deduped, instead + of citing nothing. +""" + +from __future__ import annotations + +import asyncio +import json +import time +from typing import Any + +from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_explore.config import ( + AGENT_EXPLORE_MAX_COMPLETION_TOKENS, + AGENT_EXPLORE_MODEL, + FINISH_TOOL_DESCRIPTION, + FINISH_TOOL_NAME, + FINISH_TOOL_SCHEMA, + LOOP_CONTRACT_SUFFIX, +) +from shared.services.retrieval.agent_explore.dispatch import DbFactory, dispatch_tool_call +from shared.services.retrieval.agent_explore.shared import ( + EVIDENCE_TOOL_NAMES, + budget_status_line, + build_wire_tool_name_map, + dedup_refs, + normalize_finish_refs, + tool_message_content, + wire_safe_tool_name, +) +from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult +from shared.services.retrieval.agent_tools import REGISTRY, ToolBudget, load_corpus_schema_text + +# A tool-role message is kept in full for the turn it was produced plus this +# many additional turns, then collapsed to a placeholder — see module +# docstring point 1. Not tuned against a real recall-vs-token tradeoff yet; +# 2 was chosen so a result stays fully visible for one full turn after the +# one it was produced in (enough for the model to act on it immediately), +# revisit with more Phase 4 data. +_TOOL_MESSAGE_FRESH_TURNS = 2 + + +def _resolve_client_and_model() -> tuple[Any, str]: + """Resolve the OpenAI-compatible client and ``AGENT_EXPLORE_MODEL``.""" + from shared.services.ai.llm_overrides import resolve_text + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + requested = AGENT_EXPLORE_MODEL + effective_model, api_key, api_url = resolve_text(requested) + model = effective_model or requested + client = get_openai_client(model=model, api_key=api_key, api_url=api_url) + return client, model + + +def _build_openai_tools() -> tuple[list[dict[str, Any]], dict[str, str]]: + """Return ``(tools, name_map)`` where ``name_map`` maps the wire-safe + name back to the canonical ``REGISTRY`` name (``finish`` maps to itself). + """ + specs = REGISTRY.all() + name_map = build_wire_tool_name_map([spec.name for spec in specs]) + name_map[FINISH_TOOL_NAME] = FINISH_TOOL_NAME + tools: list[dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": wire_safe_tool_name(spec.name), + "description": spec.description, + "parameters": spec.json_schema, + }, + } + for spec in specs + ] + tools.append( + { + "type": "function", + "function": { + "name": FINISH_TOOL_NAME, + "description": FINISH_TOOL_DESCRIPTION, + "parameters": FINISH_TOOL_SCHEMA, + }, + } + ) + return tools, name_map + + +def _safe_json_loads(raw: str | None) -> dict[str, Any]: + if not raw: + return {} + try: + parsed = json.loads(raw) + except (TypeError, ValueError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _collapse_stale_tool_messages( + messages: list[dict[str, Any]], + tool_message_log: list[dict[str, Any]], + *, + current_turn: int, + fresh_turns: int, +) -> None: + """Replace tool messages older than ``fresh_turns`` with a placeholder. + + ``messages`` only ever grows within one episode (see module docstring + point 1); this is what keeps that growth bounded instead of resending + every past tool result on every later turn. + """ + for entry in tool_message_log: + if entry["collapsed"]: + continue + if current_turn - entry["turn_index"] < fresh_turns: + continue + messages[entry["message_index"]]["content"] = ( + f"[collapsed: {entry['tool_name']} result from turn " + f"{entry['turn_index']} was {entry['original_chars']} chars — " + "call the tool again if you need it back in view]" + ) + entry["collapsed"] = True + + +class OpenAIHarness: + """``Harness`` implementation over an OpenAI-compatible function-calling loop.""" + + async def run_episode( + self, + *, + db_factory: DbFactory, + user_id: str, + namespace: str, + query: str, + budget: EpisodeBudget, + ) -> EpisodeResult: + tool_budget = ToolBudget() + client, model = _resolve_client_and_model() + openai_tools, tool_name_map = _build_openai_tools() + + system_prompt = load_corpus_schema_text() + LOOP_CONTRACT_SUFFIX + messages: list[dict[str, Any]] = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": query}, + ] + + steps: list[AgentStep] = [] + result_refs: list[dict[str, Any]] = [] + result_notes = "" + # Refs from every corpus.read/corpus.assets call this episode, in call + # order — the fallback source when finish's own refs end up empty (see + # module docstring point 2). + trajectory_refs: list[dict[str, Any]] = [] + # One entry per appended tool-role message: {message_index, turn_index, + # tool_name, original_chars, collapsed} — see _collapse_stale_tool_messages. + tool_message_log: list[dict[str, Any]] = [] + turn_index = 0 + + while True: + turn_index += 1 + _collapse_stale_tool_messages( + messages, + tool_message_log, + current_turn=turn_index, + fresh_turns=_TOOL_MESSAGE_FRESH_TURNS, + ) + + forced_reason = budget.exhausted() + tool_choice: Any = "auto" + if forced_reason is not None: + tool_choice = {"type": "function", "function": {"name": FINISH_TOOL_NAME}} + + turn_started = time.perf_counter() + response, usage = await asyncio.to_thread( + client.chat_completion_raw_with_usage, + messages=messages, + model=model, + temperature=0.0, + max_tokens=AGENT_EXPLORE_MAX_COMPLETION_TOKENS, + tools=openai_tools, + tool_choice=tool_choice, + ) + budget.record_usage(usage) + budget.record_step() + turn_elapsed_ms = int((time.perf_counter() - turn_started) * 1000) + turn_tokens = int((usage or {}).get("total_tokens", 0) or 0) + + message = response.choices[0].message + tool_calls = list(message.tool_calls or []) + + if not tool_calls: + stop_reason = f"budget_{forced_reason}" if forced_reason else "no_tool_call" + result_notes = str(message.content or "") + steps.append( + AgentStep( + step_index=len(steps), + tool_name="", + tool_args={}, + observation_text=result_notes, + error=None, + elapsed_ms=turn_elapsed_ms, + tokens_used_delta=turn_tokens, + tokens_used_total=budget.tokens_used, + ) + ) + break + + finish_call = next( + (tc for tc in tool_calls if tc.function.name == FINISH_TOOL_NAME), None + ) + if finish_call is not None: + args = _safe_json_loads(finish_call.function.arguments) + result_refs = normalize_finish_refs(args.get("refs")) + result_notes = str(args.get("notes") or "") + stop_reason = f"budget_{forced_reason}" if forced_reason else "finished" + steps.append( + AgentStep( + step_index=len(steps), + tool_name=FINISH_TOOL_NAME, + tool_args=args, + observation_text=f"refs={len(result_refs)} notes={result_notes!r}", + error=None, + elapsed_ms=turn_elapsed_ms, + tokens_used_delta=turn_tokens, + tokens_used_total=budget.tokens_used, + ) + ) + break + + if forced_reason is not None: + # Forced tool_choice=finish but the provider returned a + # different tool anyway (not observed in verification, but a + # budget cutoff must never loop past). Stop here regardless. + stop_reason = f"budget_{forced_reason}" + result_notes = str(message.content or "") or ( + "budget exhausted; provider did not return finish" + ) + steps.append( + AgentStep( + step_index=len(steps), + tool_name="", + tool_args={}, + observation_text=result_notes, + error="forced_finish_not_honored", + elapsed_ms=turn_elapsed_ms, + tokens_used_delta=turn_tokens, + tokens_used_total=budget.tokens_used, + ) + ) + break + + messages.append( + { + "role": "assistant", + "content": message.content or "", + "tool_calls": [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in tool_calls + ], + } + ) + first_tool_tokens_recorded = False + for tc in tool_calls: + tool_started = time.perf_counter() + args = _safe_json_loads(tc.function.arguments) + requested_name = str(tc.function.name or "") + canonical_name = tool_name_map.get(requested_name, requested_name) + tool_result = await dispatch_tool_call( + canonical_name, + args, + db_factory=db_factory, + user_id=user_id, + namespace=namespace, + budget=tool_budget, + ) + tool_elapsed_ms = int((time.perf_counter() - tool_started) * 1000) + content = tool_message_content(tool_result, max_chars=tool_budget.max_chars) + messages.append( + {"role": "tool", "tool_call_id": tc.id, "content": content} + ) + tool_message_log.append( + { + "message_index": len(messages) - 1, + "turn_index": turn_index, + "tool_name": canonical_name, + "original_chars": len(content), + "collapsed": False, + } + ) + if canonical_name in EVIDENCE_TOOL_NAMES and not tool_result.error: + trajectory_refs.extend(tool_result.refs) + # Turn-level token usage is attributed to the first tool step in + # this turn (the completion that decided all calls in it); the + # rest are 0 to avoid double-counting the same LLM usage. + steps.append( + AgentStep( + step_index=len(steps), + tool_name=canonical_name, + tool_args=args, + observation_text=content, + error=tool_result.error, + elapsed_ms=( + tool_elapsed_ms + if first_tool_tokens_recorded + else turn_elapsed_ms + tool_elapsed_ms + ), + tokens_used_delta=0 if first_tool_tokens_recorded else turn_tokens, + tokens_used_total=budget.tokens_used, + ) + ) + first_tool_tokens_recorded = True + + # Appended once per turn, to the last tool message only (not + # every AgentStep's recorded observation_text above) — the model + # only needs to see current remaining budget once before its next + # completion call, not once per parallel tool call in this turn. + messages[-1]["content"] = ( + str(messages[-1]["content"]) + "\n" + budget_status_line(budget) + ) + + if not result_refs: + fallback_refs = dedup_refs(trajectory_refs) + if fallback_refs: + result_refs = fallback_refs + result_notes = (result_notes + " " if result_notes else "") + ( + "[refs auto-filled from corpus.read/corpus.assets trajectory; " + "finish did not cite any]" + ) + + return EpisodeResult( + refs=result_refs, + notes=result_notes, + steps=steps, + stop_reason=stop_reason, + tokens_used=budget.tokens_used, + model_name=model, + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/resolve.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/resolve.py new file mode 100644 index 00000000..cf19cd86 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/resolve.py @@ -0,0 +1,40 @@ +"""``AGENT_EXPLORE_HARNESS`` env switch: unrecognized or unset values fall +back to the default rather than raising, so a typo'd env var degrades to +known-good behavior instead of breaking the route. + +Each branch below imports its harness implementation lazily so that +selecting ``openai`` never imports ``cursor_sdk``-dependent code (and +vice versa) — see ``cursor_harness.py``'s guarded import. +""" + +from __future__ import annotations + +import os + +from shared.services.retrieval.agent_explore.harness.base import Harness + +_HARNESS_ENV = "AGENT_EXPLORE_HARNESS" +_HARNESSES = {"openai", "cursor_sdk"} +_DEFAULT_HARNESS = "cursor_sdk" + + +def resolve_harness_name() -> str: + """``cursor_sdk`` (default) or ``openai``.""" + value = os.environ.get(_HARNESS_ENV, "").strip().lower() + return value if value in _HARNESSES else _DEFAULT_HARNESS + + +def resolve_harness(name: str | None = None) -> Harness: + """Build the ``Harness`` implementation for ``name`` (default: env-resolved).""" + resolved = (name or resolve_harness_name()).strip().lower() + if resolved == "cursor_sdk": + from shared.services.retrieval.agent_explore.harness.cursor_harness import ( + CursorHarness, + ) + + return CursorHarness() + from shared.services.retrieval.agent_explore.harness.openai_harness import ( + OpenAIHarness, + ) + + return OpenAIHarness() diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py b/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py new file mode 100644 index 00000000..5edeeac4 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py @@ -0,0 +1,106 @@ +"""Resolve ``finish.refs`` into the chunk_id-bearing shape ``resolve_workflow_references`` requires. + +Verified live: ``resolve_workflow_references`` -> ``hydrate_referenced_chunk_rows`` +drops any ref whose lookup key has an empty ``chunk_id`` (``row_utils. +build_reference_lookup_key`` + the ``ref_keys = [k for k in ref_keys if k[0] +and k[1]]`` guard in ``hydration/reference.py``) — a ``{document_id, +section_path}``-only ref silently resolves to zero ``referenced_chunks``, +which is exactly the shape ``agent_tools.CORPUS_SCHEMA.md``/``corpus.read`` +teaches the agent to cite (``corpus.read``'s rendered ``text`` — the only +thing the LLM ever sees — shows ``section_path``, never ``chunk_id``; +``chunk_id`` only appears in its structured ``payload``/``refs``, which the +LLM does not see). This module closes that gap at the harness boundary +instead of changing what the agent is taught to cite: for any ref missing +``chunk_id``, resolve that section's own body chunk — the same "one section, +one body chunk" lookup ``corpus.read``'s ``section_path`` branch performs +(``agent_tools/section_path_lookup.py``), including the same suffix fallback +when the agent cites a path without ancestor prefixes. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.section_path_lookup import ( + resolve_section_path_anchor, +) + +_BODY_CHUNK_TYPES = ("text", "page") + + +async def resolve_finish_refs( + db: AsyncSession, + *, + user_id: str, + namespace: str, + refs: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Return refs with ``chunk_id`` populated; drops refs that don't resolve.""" + document_ids = { + str(ref.get("document_id") or "").strip() for ref in refs if ref.get("document_id") + } + if not document_ids: + return [] + + documents = ( + ( + await db.execute( + select(Document) + .where(Document.document_id.in_(document_ids)) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + ) + ) + .scalars() + .all() + ) + revision_by_doc = { + d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + } + + resolved: list[dict[str, Any]] = [] + for ref in refs: + document_id = str(ref.get("document_id") or "").strip() + chunk_id = str(ref.get("chunk_id") or "").strip() + if document_id and chunk_id: + resolved.append({"document_id": document_id, "chunk_id": chunk_id}) + continue + + section_path = str(ref.get("section_path") or "").strip() + job_result_id = revision_by_doc.get(document_id) + if not (document_id and section_path and job_result_id): + continue + + resolved_path, path_error = await resolve_section_path_anchor( + db, + document_id=document_id, + job_result_id=job_result_id, + section_path=section_path, + ) + if path_error or not resolved_path: + continue + + row = ( + await db.execute( + select(DocumentChunk.chunk_id) + .select_from(DocumentChunk) + .join( + DocumentSection, + DocumentSection.section_id == DocumentChunk.section_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentSection.section_path == resolved_path) + .where(DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES)) + ) + ).first() + if row is None: + continue + resolved.append({"document_id": document_id, "chunk_id": str(row[0])}) + + return resolved diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/shared.py b/packages/shared-python/shared/services/retrieval/agent_explore/shared.py new file mode 100644 index 00000000..36223c4d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/shared.py @@ -0,0 +1,126 @@ +"""Provider-agnostic helpers shared by every ``agent_explore`` harness. + +Extracted from ``episode.py`` (originally OpenAI-harness-specific) once a +second harness (``harness/cursor_harness.py``) needed the exact same logic +and, as a debug-script PoC, had started duplicating and drifting from it +instead of sharing it — see the Phase 3.5 section of +``.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md``. + +Deliberately excludes anything that assumes an editable ``messages: +list[dict]`` conversation history (that's OpenAI-harness-specific — see +``harness/openai_harness.py``'s ``_collapse_stale_tool_messages``, which is +NOT here because the Cursor SDK manages its own context with no equivalent +hook exposed to the host process). +""" + +from __future__ import annotations + +from typing import Any + +from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_tools import ToolResult + +# Tools whose ToolResult.refs point at evidence the agent has actually looked +# at (full body content), as opposed to candidate/listing refs from +# list_documents/outline/node_filter/recall/grep — those describe *where +# things are*, not *what was read*, and would inject unread noise into the +# trajectory-refs fallback below if included. +EVIDENCE_TOOL_NAMES = frozenset({"corpus.read", "corpus.assets"}) + + +def wire_safe_tool_name(name: str) -> str: + """Replace ``.`` with ``_`` in a tool name for function-calling wire formats. + + Every ``agent_tools`` name is dotted (``corpus.read``); DeepSeek's + (OpenAI-compatible) function-calling API rejects ``.`` in + ``tools[].function.name`` (must match ``^[a-zA-Z0-9_-]+$``, verified + live), and the Cursor SDK PoC needed the identical replacement for its + own ``custom_tools`` wire names — this is a function-calling wire-format + restriction shared by both providers, not an OpenAI-specific quirk. The + dotted name stays canonical in ``REGISTRY``/MCP; this underscore form + exists only for providers whose wire format rejects dots. + """ + return name.replace(".", "_") + + +def build_wire_tool_name_map(names: list[str]) -> dict[str, str]: + """``{wire_safe_name: canonical_name}`` for every name in ``names``. + + Names that are already wire-safe (e.g. ``finish``, which has no dot) map + to themselves. Used by a harness to translate a provider's tool-call + name back to the canonical ``REGISTRY`` name before dispatch. + """ + return {wire_safe_tool_name(name): name for name in names} + + +def tool_message_content(result: ToolResult, *, max_chars: int) -> str: + """Cap a tool's rendered text before it enters LLM context. + + Uses the caller's ``ToolBudget.max_chars`` (``EVIDENCE_TEXT_CHAR_BUDGET``, + aligned with evidence packing — see ``agent_tools/registry.py``) + so tools like ``read`` can return unbounded body text while the harness + still bounds what the model sees per turn. This cap applies uniformly to + every tool's rendered text (not just ``read``'s body content) — a tool + that returns a "complete, non-truncated" *matched set* by contract + (``outline``, ``node_filter``) still has its *rendered text* capped here + the same as any other tool; that promise is about payload/refs + cardinality, not about how much of it is shown to the LLM per turn. + """ + if result.error: + return f"error: {result.error}" + text = result.text or "(empty result)" + if len(text) <= max_chars: + return text + omitted = len(text) - max_chars + return ( + text[:max_chars] + + f"\n...[truncated, {omitted} more chars — narrow the scope " + "(e.g. depth/path_prefix for outline, a tighter predicate for " + "node_filter, or a more specific ref for read) and call again if " + "you need the rest]" + ) + + +def normalize_finish_refs(raw: Any) -> list[dict[str, Any]]: + """Keep only dict items with a non-empty ``document_id`` from a raw ``finish.refs``.""" + if not isinstance(raw, list): + return [] + normalized: list[dict[str, Any]] = [] + for item in raw: + if isinstance(item, dict) and str(item.get("document_id") or "").strip(): + normalized.append(item) + return normalized + + +def budget_status_line(budget: EpisodeBudget) -> str: + """One-line remaining-budget summary appended to a tool observation. + + Neither harness previously surfaced ``EpisodeBudget``'s own counters + (``steps_used``/``max_steps``, ``tokens_used``/``token_limit``, + elapsed/wall_clock) to the model at all — it had no way to tell "I'm on + step 3 of 12" from "I'm on step 11 of 12", so it could not self-regulate + when to stop exploring and call ``finish``. This exposes the same + ``EpisodeBudget.snapshot()`` data already used for the hard cutoff, + reused as a soft signal the model can read every turn. + """ + snap = budget.snapshot() + return ( + f"[budget: steps {snap['steps_used']}/{snap['max_steps']}, " + f"tokens {snap['tokens_used']}/{snap['token_limit']}, " + f"elapsed {snap['elapsed_seconds']:.0f}s/{snap['wall_clock_seconds']:.0f}s]" + ) + + +def dedup_refs(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Dedup by ``(document_id, chunk_id)``, keeping first-seen order.""" + seen: set[tuple[str, str]] = set() + deduped: list[dict[str, Any]] = [] + for ref in refs: + document_id = str(ref.get("document_id") or "").strip() + chunk_id = str(ref.get("chunk_id") or "").strip() + key = (document_id, chunk_id) + if not document_id or not chunk_id or key in seen: + continue + seen.add(key) + deduped.append(ref) + return deduped diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/types.py b/packages/shared-python/shared/services/retrieval/agent_explore/types.py new file mode 100644 index 00000000..89bc90da --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/types.py @@ -0,0 +1,32 @@ +"""Result types shared between ``episode.py`` and ``bridge.py``.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class AgentStep: + """One tool call executed during the episode (one LLM turn may yield several).""" + + step_index: int + tool_name: str + tool_args: dict[str, Any] + observation_text: str + error: str | None + elapsed_ms: int + tokens_used_delta: int + tokens_used_total: int + + +@dataclass +class EpisodeResult: + """Everything ``bridge.py`` / the route need after the episode ends.""" + + refs: list[dict[str, Any]] + notes: str + steps: list[AgentStep] = field(default_factory=list) + stop_reason: str = "finished" + tokens_used: int = 0 + model_name: str = "" diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md index 496b1607..ec65d189 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md +++ b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md @@ -54,6 +54,14 @@ they differ in shape, and in **which sections get a body chunk at all**: body chunk. Internal/structural sections never carry a `page` chunk themselves — their summaries aggregate from their leaf descendants. +A section owns **at most one** body chunk either way (verified against +real published data: a section with children and a section without both +had exactly 0 or 1, never more). So a "chunk" and "the section that owns +it" are the same unit, not two different granularities — there is no +separate finer-or-coarser level to choose between. The only structural +sections with no chunk of their own are `chunk`-track sections whose +entire content lives in their descendants. + For `page` chunks specifically: one leaf section's body may span one or more physical pages. A page's text is stored **once**, under whichever leaf is first in reading order to cover that page (the "owner"). Every other @@ -115,7 +123,9 @@ for anything finer-grained than a document pair. ## 5. Reserved / not yet available - **Vector channel**: `recall`'s `channels` parameter reserves a `vector` - option; it does not exist yet. `recall` today is lexical only. + option; it does not exist yet. `recall` today fuses two lexical channels + (`path_content`: persisted map-unit BM25 over path+content; `term`: + substring match over `document_map_units.term_search_text_lower`) via RRF. ## 6. Tools and when to use each @@ -125,8 +135,8 @@ for anything finer-grained than a document pair. | `outline` | The task only needs titles/summaries — overview, "what does chapter N cover," picking where to look before reading | one document, or a `section_path` prefix within it | Titles + summaries + `chunk_count`, no body text, no folding. Depth-limited by argument, not by a token budget. Use this to build your own map instead of relying on a pre-folded one. | | `node_filter` | The task is a traversal/exclusion predicate — FOR ALL / EXISTS / ANY / NOT — over section titles or summaries ("which docs mention X in a heading," "sections NOT about Y") | one or more documents | Deterministic substring/regex match against `section_path` and `summary` only, not body text. Returns the full matching set and count, never a truncated top-K. If the predicate must run against body text, use `grep` instead. | | `grep` | Exact string / regex / identifier / number lookup that must run against body text | scoped by document/section/chunk_type | Returns match count plus snippets, so ANY/ALL logic can also close over body text, not just titles. | -| `recall` | A fuzzy question where you don't know where the answer lives | namespace or scoped | Ranked candidates (path + content BM25 today; term and vector are separate/reserved — see §5) with path and snippet, not full content. | -| `read` | You already know which section(s)/chunk(s) to read | one or more sections/chunks | Returns full body content, resolves `SAME-AS` markers into the owner's text, expands `connect_to` assets, and converts `page_assets` into URLs. | +| `recall` | A fuzzy question where you don't know where the answer lives | namespace or scoped | Ranked candidates from `path_content` (BM25) + `term` (substring) channels fused by RRF; `vector` is reserved — see §5. Returns path and snippet, not full content. | +| `read` | You already know which section(s)/chunk(s) to read | one or more sections/chunks | Returns full body content, resolves `SAME-AS` markers into the owner's text, expands `connect_to` assets, and converts `page_assets` into URLs. If your `section_path` omits ancestor segments (e.g. missing a top-level volume like `附件目录 /`), `read` tries a unique suffix match within the document; if several sections match, it returns an ambiguity error listing the full paths — copy the full path from `outline`/`grep`/`refs` when that happens. | | `assets` | You need images/tables directly, or need to find which section(s) host a given asset | one or more documents | Forward (by type/query) and reverse (asset → hosting section) lookup — see §3. | | `neighbors` | You need related documents in the same namespace | one document | Document-level `related` edges only — see §4. | diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py b/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py new file mode 100644 index 00000000..5583243f --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py @@ -0,0 +1,51 @@ +"""Provider-agnostic corpus exploration tools. + +See ``.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md`` (Phase 2) +for the design. Tools in this package query the published, DB-served corpus +described in ``CORPUS_SCHEMA.md`` (``documents`` / ``document_sections`` / +``document_chunks`` / ``graph_nodes`` / ``graph_edges``) — not the on-disk +parse artifacts. + +The same ``REGISTRY`` is meant to be consumed by two harnesses (Phase 3): +the API ``/mcp`` server (Cursor/Codex/Claude) and the in-process +``agent_explore`` tool-loop. Importing ``agent_tools.tools`` registers every +tool as a side effect; this package performs that import itself (below), so +any consumer importing a name here (e.g. ``REGISTRY``) is guaranteed a fully +populated registry. +""" + +from __future__ import annotations + +from shared.services.retrieval.agent_tools.registry import ( + REGISTRY, + ToolBudget, + ToolContext, + ToolRegistry, + ToolResult, + ToolSpec, + capped_limit, + register_tool, +) +from shared.services.retrieval.agent_tools.schema_doc import load_corpus_schema_text + +# Register every ``corpus.*`` tool into ``REGISTRY`` as an import side effect: +# each module under ``tools/`` calls ``@register_tool`` at import time. Kept +# here — not in ``registry.py``, which the tool modules import and would +# therefore create a cycle — so consumers importing any name from this package +# no longer need their own side-effect import of ``tools``. ``_tools`` is +# intentionally never referenced (listed in ``__all__`` as a deliberate +# re-export so tooling does not flag it as unused). +from shared.services.retrieval.agent_tools import tools as _tools + +__all__ = [ + "_tools", + "REGISTRY", + "ToolBudget", + "ToolContext", + "ToolRegistry", + "ToolResult", + "ToolSpec", + "capped_limit", + "load_corpus_schema_text", + "register_tool", +] diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py new file mode 100644 index 00000000..d5aa871c --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py @@ -0,0 +1,142 @@ +"""Provider-agnostic tool contracts for corpus exploration. + +Mirrors the shape of ``apps/worker/app/services/document_agent/registry.py`` +(``ToolSpec`` + a decorator-based registry), adapted for the async DB-backed +corpus tools in this package: ``ToolSpec(name, description, json_schema, run)``, +``ToolContext(db, user_id, namespace, budget)``, ``ToolResult(text, payload, refs)``. + +Both the API ``/mcp`` server and the in-process ``agent_explore`` tool-loop +(Phase 3) dispatch through the same ``REGISTRY`` — this module has no +provider-specific (MCP / OpenAI tool-calling) concerns. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.settings import EVIDENCE_TEXT_CHAR_BUDGET + + +@dataclass(frozen=True) +class ToolBudget: + """Per-call output budget passed down to every tool via ``ToolContext``. + + ``max_items`` is a hard ceiling on how many rows a tool that returns an + unbounded/ranked list (``recall``, ``grep``, asset forward search) may + return in one call: each such tool keeps its own smaller, tool-appropriate + default (e.g. ``grep``'s ``max_results``, ``recall``'s ``top_k``) but + clamps the caller-requested value to this ceiling via ``capped_limit()`` + below, and notes it in ``ToolResult.text`` when the request was clamped. + Tools that promise a complete, non-truncated set by contract + (``node_filter``, ``outline``) do not apply this budget to their + matched-set cardinality — see ``CORPUS_SCHEMA.md`` §6. + + ``max_chars`` caps the rendered ``ToolResult.text`` before it enters LLM + context. Applied in ``agent_explore.shared.tool_message_content`` (not + inside individual tools) so ``read`` can return full body text from the + tool while the harness still bounds what the model sees per turn. Aligned + with final evidence packing via ``EVIDENCE_TEXT_CHAR_BUDGET`` + (12_000). + """ + + max_chars: int = EVIDENCE_TEXT_CHAR_BUDGET + max_items: int = 50 + + +def capped_limit(requested: int, budget: ToolBudget) -> int: + """Clamp a caller-requested row count to ``budget.max_items`` (min 1).""" + return max(1, min(requested, budget.max_items)) + + +@dataclass +class ToolContext: + """Per-call execution context. One instance is built per tool dispatch.""" + + db: AsyncSession + user_id: str + namespace: str + budget: ToolBudget = field(default_factory=ToolBudget) + + +@dataclass +class ToolResult: + """Uniform tool output. + + ``text`` is the human/LLM-facing rendering; ``payload`` is the structured + data (for programmatic callers and for building ``refs``); ``refs`` are + resolvable evidence pointers (``{document_id, section_path|chunk_id}``) + that a harness can fold into ``referenced_chunks`` (Phase 3 bridge). + ``error`` is set instead of raising for caller-facing input mistakes (bad + args, unknown document_id) so a tool-loop agent can see and correct them. + """ + + text: str + payload: dict[str, Any] = field(default_factory=dict) + refs: list[dict[str, Any]] = field(default_factory=list) + error: str | None = None + + +ToolHandler = Callable[[ToolContext, dict[str, Any]], Awaitable[ToolResult]] + + +@dataclass(frozen=True) +class ToolSpec: + name: str + description: str + json_schema: dict[str, Any] + run: ToolHandler + + +class ToolRegistry: + """Name -> ``ToolSpec`` map. Provider-agnostic; no MCP/OpenAI coupling.""" + + def __init__(self) -> None: + self._tools: dict[str, ToolSpec] = {} + + def register(self, spec: ToolSpec) -> None: + if spec.name in self._tools: + raise ValueError(f"tool already registered: {spec.name}") + self._tools[spec.name] = spec + + def get(self, name: str) -> ToolSpec | None: + return self._tools.get(name) + + def all(self) -> list[ToolSpec]: + return list(self._tools.values()) + + async def dispatch( + self, name: str, ctx: ToolContext, args: dict[str, Any] + ) -> ToolResult: + spec = self.get(name) + if spec is None: + return ToolResult(text="", error=f"unknown tool: {name}") + return await spec.run(ctx, args) + + +REGISTRY = ToolRegistry() + + +def register_tool( + *, + name: str, + description: str, + json_schema: dict[str, Any], +) -> Callable[[ToolHandler], ToolHandler]: + """Decorator mirroring worker's ``register_tool`` for the async corpus tools.""" + + def _decorator(handler: ToolHandler) -> ToolHandler: + REGISTRY.register( + ToolSpec( + name=name, + description=description, + json_schema=json_schema, + run=handler, + ) + ) + return handler + + return _decorator diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py b/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py new file mode 100644 index 00000000..37569589 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py @@ -0,0 +1,21 @@ +"""Single loader for ``CORPUS_SCHEMA.md`` — the agent-facing corpus schema text. + +Both harnesses (Phase 3) read through this function instead of the file +directly, so there is exactly one place that resolves the path: the API +``/mcp`` server's ``instructions`` and ``agent_explore``'s system prompt must +stay byte-identical for the shared schema portion (see the module docstring +at the top of ``CORPUS_SCHEMA.md`` — "do not duplicate it elsewhere"). +""" + +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path + +_SCHEMA_PATH = Path(__file__).with_name("CORPUS_SCHEMA.md") + + +@lru_cache(maxsize=1) +def load_corpus_schema_text() -> str: + """Return the verbatim contents of ``CORPUS_SCHEMA.md``.""" + return _SCHEMA_PATH.read_text(encoding="utf-8") diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py b/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py new file mode 100644 index 00000000..3df35992 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py @@ -0,0 +1,87 @@ +"""Shared section_path resolution for agent tools and harness bridge. + +``corpus.read`` and ``resolve_finish_refs`` both need to turn an agent-supplied +``section_path`` into one canonical DB path. Agents often cite a suffix (e.g. +``3 工程地质 / 3.2 覆盖层``) while the stored path includes ancestors +(``附件目录 / 3 工程地质 / 3.2 覆盖层``). Exact match alone fails silently +downstream; this module adds a segment-bound suffix fallback and surfaces +ambiguity instead of guessing. +""" + +from __future__ import annotations + +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import DocumentSection +from shared.services.retrieval.search.lexical_text import normalize_section_path + + +def paths_matching_section_ref(normalized: str, candidate_paths: list[str]) -> list[str]: + """Return paths that equal ``normalized`` or end with `` / {normalized}``.""" + if not normalized: + return [] + suffix = f" / {normalized}" + matches = [ + path + for path in candidate_paths + if path == normalized or (normalized != "Root" and path.endswith(suffix)) + ] + return sorted(dict.fromkeys(matches)) + + +def format_ambiguous_section_path_error(normalized: str, matches: list[str]) -> str: + return ( + f"ambiguous section_path {normalized!r}: matches {len(matches)} sections — " + "use the full path from outline/grep/refs: " + + "; ".join(matches) + ) + + +def section_path_anchor_filter(resolved_path: str): + """SQL filter for one section (``mode=self`` anchor).""" + return DocumentSection.section_path == resolved_path + + +def section_path_subtree_filter(resolved_path: str): + """SQL filter for a section and all descendants (``mode=descendants``).""" + return or_( + DocumentSection.section_path == resolved_path, + DocumentSection.section_path.like(f"{resolved_path} / %"), + ) + + +async def resolve_section_path_anchor( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + section_path: str, +) -> tuple[str | None, str | None]: + """Resolve one canonical ``section_path`` or return ``(None, error)``.""" + normalized = normalize_section_path(section_path) + base = ( + select(DocumentSection.section_path) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + ) + + exact_row = (await db.execute(base.where(DocumentSection.section_path == normalized))).first() + if exact_row is not None and exact_row[0]: + return str(exact_row[0]), None + + suffix_filter = or_( + DocumentSection.section_path == normalized, + DocumentSection.section_path.like(f"% / {normalized}"), + ) + suffix_rows = ( + await db.execute(base.where(suffix_filter).order_by(DocumentSection.sort_order)) + ).all() + matches = paths_matching_section_ref( + normalized, [str(row[0]) for row in suffix_rows if row and row[0]] + ) + if not matches: + return None, f"unknown section_path for {document_id}: {normalized}" + if len(matches) > 1: + return None, format_ambiguous_section_path_error(normalized, matches) + return matches[0], None diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py b/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py new file mode 100644 index 00000000..6dfed42c --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py @@ -0,0 +1,72 @@ +"""Shared hit-anchored snippet builder for ``corpus.grep`` and ``corpus.recall``. + +Both tools locate a single pattern/term inside one chunk's text and need to +show a bounded excerpt around it. The shape is: head anchor + the window +around the (first) match + tail anchor, with ``...`` between spans that do +not touch. Overlapping/adjacent spans are merged before rendering so short +chunks never produce duplicate text or a stray ``...`` inside otherwise +continuous text. + +Only the first match is windowed. A single call passes a single +pattern/term, but that pattern can still occur more than once inside one +chunk (verified against real data); later occurrences in the same chunk are +not separately windowed here — use ``corpus.read`` on the chunk for the rest. +""" + +from __future__ import annotations + +HIT_CONTEXT_CHARS = 80 +HEAD_TAIL_CHARS = 50 + + +def _merge_spans(spans: list[tuple[int, int]]) -> list[tuple[int, int]]: + ordered = sorted(s for s in spans if s[1] > s[0]) + merged: list[list[int]] = [] + for start, end in ordered: + if merged and start <= merged[-1][1]: + merged[-1][1] = max(merged[-1][1], end) + else: + merged.append([start, end]) + return [(start, end) for start, end in merged] + + +def build_snippet( + text: str, + hit: tuple[int, int] | None = None, + *, + hit_context: int = HIT_CONTEXT_CHARS, + head_tail: int = HEAD_TAIL_CHARS, +) -> str: + """Head/tail anchor + first-match window, joined by ``...`` where spans don't touch. + + ``hit`` is the ``(start, end)`` char offset of the located match in + ``text``, or ``None`` when no specific position is known (falls back to + head/tail anchors only). Short text (<= ``head_tail * 2`` chars) is + returned unchanged. + """ + if not text: + return "" + if len(text) <= head_tail * 2: + return text + + spans: list[tuple[int, int]] = [ + (0, head_tail), + (max(len(text) - head_tail, 0), len(text)), + ] + if hit is not None: + hit_start, hit_end = hit + spans.append( + (max(hit_start - hit_context, 0), min(hit_end + hit_context, len(text))) + ) + + merged = _merge_spans(spans) + parts: list[str] = [] + prev_end = 0 + for start, end in merged: + if start > prev_end: + parts.append("...") + parts.append(text[start:end]) + prev_end = end + if prev_end < len(text): + parts.append("...") + return "".join(parts) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/__init__.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/__init__.py new file mode 100644 index 00000000..f94f447e --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/__init__.py @@ -0,0 +1,29 @@ +"""Importing this package registers every ``corpus.*`` tool into ``REGISTRY``. + +One module per tool, mirroring +``apps/worker/app/services/document_agent/tools/``. +""" + +from __future__ import annotations + +from shared.services.retrieval.agent_tools.tools import ( + assets as _assets, + grep as _grep, + list_documents as _list_documents, + neighbors as _neighbors, + node_filter as _node_filter, + outline as _outline, + read as _read, + recall as _recall, +) + +__all__ = [ + "_assets", + "_grep", + "_list_documents", + "_neighbors", + "_node_filter", + "_outline", + "_read", + "_recall", +] diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py new file mode 100644 index 00000000..18c5a1b6 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py @@ -0,0 +1,193 @@ +"""``corpus.assets`` — forward asset search and reverse asset -> hosts lookup. + +Image/table chunks are parked under their document's synthetic ``Root`` +section in the DB (§3 of ``CORPUS_SCHEMA.md``); the real association to a +body section lives in ``chunk_metadata.connect_to`` on the *body* chunk, not +on the asset. There is no stored asset -> body back-link, so the reverse +lookup (``host_of``) scans the candidate documents' text/page chunks in +Python and checks ``connect_to`` for the requested target ids. + +``chunk_metadata`` is a plain ``JSON`` column (not ``JSONB``), so a +containment query (``@>``) is not available here — that operator is +JSONB-only in PostgreSQL. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + register_tool, +) +from shared.services.retrieval.hydration.row_utils import iter_connected_target_ids +from shared.services.retrieval.settings import ASSET_CHUNK_TYPES + +_BODY_CHUNK_TYPES = ("text", "page") + + +@register_tool( + name="corpus.assets", + description=( + "Forward search for image/table chunks by type/query, or reverse " + "lookup: given asset chunk_ids (host_of), find which body " + "section(s) embed or reference them via connect_to." + ), + json_schema={ + "type": "object", + "properties": { + "document_ids": {"type": "array", "items": {"type": "string"}}, + "type": { + "type": "string", + "enum": ["image", "table", "any"], + "default": "any", + }, + "query": { + "type": "string", + "description": "Substring match against summary/keywords (forward search only).", + }, + "host_of": { + "type": "array", + "items": {"type": "string"}, + "description": "Asset chunk_ids to reverse-resolve to hosting sections.", + }, + }, + "required": [], + }, +) +async def assets(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + document_ids = [ + str(d).strip() for d in (args.get("document_ids") or []) if str(d).strip() + ] + host_of = [str(c).strip() for c in (args.get("host_of") or []) if str(c).strip()] + + scope_filters: list[Any] = [ + Document.user_id == ctx.user_id, + Document.namespace == ctx.namespace, + Document.status == "active", + Document.current_job_result_id == DocumentChunk.job_result_id, + ] + if document_ids: + scope_filters.append(Document.document_id.in_(document_ids)) + + if host_of: + return await _reverse_lookup(ctx, scope_filters=scope_filters, target_ids=host_of) + return await _forward_search( + ctx, + scope_filters=scope_filters, + asset_type=str(args.get("type") or "any").strip().lower(), + query=str(args.get("query") or "").strip().lower(), + ) + + +async def _forward_search( + ctx: ToolContext, + *, + scope_filters: list[Any], + asset_type: str, + query: str, +) -> ToolResult: + types = {asset_type} if asset_type in ASSET_CHUNK_TYPES else set(ASSET_CHUNK_TYPES) + stmt = ( + select(DocumentChunk, DocumentSection.section_path, Document.source_file_name) + .select_from(DocumentChunk) + .join(Document, Document.document_id == DocumentChunk.document_id) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .where(*scope_filters) + .where(DocumentChunk.chunk_type.in_(sorted(types))) + .order_by(DocumentChunk.document_id, DocumentChunk.sort_order) + ) + rows = (await ctx.db.execute(stmt)).all() + + results: list[dict[str, Any]] = [] + for chunk, section_path, source_file_name in rows: + metadata = chunk.chunk_metadata if isinstance(chunk.chunk_metadata, dict) else {} + summary = str(metadata.get("summary") or "").strip() + keywords = metadata.get("keywords") or [] + if query: + haystack = " ".join( + [summary.lower(), " ".join(str(k).lower() for k in keywords)] + ) + if query not in haystack: + continue + results.append( + { + "chunk_id": chunk.chunk_id, + "document_id": chunk.document_id, + "source_file_name": source_file_name, + "chunk_type": chunk.chunk_type, + "file_path": chunk.file_path, + "summary": summary, + "keywords": keywords, + "section_path": section_path, + } + ) + if len(results) >= ctx.budget.max_items: + break + + lines = [f"assets={len(results)}"] + for r in results: + lines.append(f"- [{r['chunk_type']}] {r['file_path']} — {r['summary']}") + + return ToolResult( + text="\n".join(lines), + payload={"assets": results}, + refs=[{"document_id": r["document_id"], "chunk_id": r["chunk_id"]} for r in results], + ) + + +async def _reverse_lookup( + ctx: ToolContext, + *, + scope_filters: list[Any], + target_ids: list[str], +) -> ToolResult: + target_set = set(target_ids) + stmt = ( + select(DocumentChunk, DocumentSection.section_path, Document.source_file_name) + .select_from(DocumentChunk) + .join(Document, Document.document_id == DocumentChunk.document_id) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .where(*scope_filters) + .where(DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES)) + ) + rows = (await ctx.db.execute(stmt)).all() + + hosts_by_target: dict[str, list[dict[str, Any]]] = {tid: [] for tid in target_set} + for chunk, section_path, source_file_name in rows: + row = {"chunk_metadata": chunk.chunk_metadata} + for target_id in iter_connected_target_ids(row): + if target_id in target_set: + hosts_by_target[target_id].append( + { + "document_id": chunk.document_id, + "source_file_name": source_file_name, + "section_path": section_path, + "chunk_id": chunk.chunk_id, + "chunk_type": chunk.chunk_type, + } + ) + + lines = [] + for target_id, hosts in hosts_by_target.items(): + if not hosts: + lines.append(f"- {target_id}: no host found (unresolved Root asset)") + continue + for host in hosts: + lines.append( + f"- {target_id} <- {host['source_file_name']} / {host['section_path']}" + ) + + return ToolResult( + text="\n".join(lines) if lines else "no hosts found", + payload={"hosts_by_target": hosts_by_target}, + refs=[ + {"document_id": host["document_id"], "section_path": host["section_path"]} + for hosts in hosts_by_target.values() + for host in hosts + ], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py new file mode 100644 index 00000000..af7224c7 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py @@ -0,0 +1,168 @@ +"""``corpus.grep`` — exact string/regex lookup against body text. + +SQL ``ILIKE`` / ``~*`` on ``document_chunks.content``, scoped to the current +revision. Reports a total match count (over the full in-scope corpus, not +just the returned page) alongside capped snippets, so ANY/ALL logic can close +over body text the same way ``corpus.node_filter`` closes over titles/summaries. + +Snippets are built by the shared ``agent_tools.snippet.build_snippet`` (head ++ first-match window + tail, ``...``-joined, overlap-merged) — the same +mechanism ``corpus.recall``'s term channel uses, so the two tools don't carry +duplicate window-slicing logic or drift to different constants. +""" + +from __future__ import annotations + +import re +from typing import Any + +from sqlalchemy import func, select + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + capped_limit, + register_tool, +) +from shared.services.retrieval.agent_tools.snippet import ( + HIT_CONTEXT_CHARS, + build_snippet, +) + +_DEFAULT_MAX_RESULTS = 30 +_DEFAULT_CONTEXT_CHARS = HIT_CONTEXT_CHARS + + +def _build_scope_filters( + *, + user_id: str, + namespace: str, + document_ids: list[str], + chunk_types: set[str], +) -> list[Any]: + filters: list[Any] = [ + Document.user_id == user_id, + Document.namespace == namespace, + Document.status == "active", + Document.current_job_result_id == DocumentChunk.job_result_id, + ] + if document_ids: + filters.append(Document.document_id.in_(document_ids)) + if chunk_types: + filters.append(func.lower(DocumentChunk.chunk_type).in_(sorted(chunk_types))) + return filters + + +@register_tool( + name="corpus.grep", + description=( + "Exact string or regex search against chunk body text (content), " + "not titles/summaries (use corpus.node_filter for that). Returns the " + "total number of matching chunks plus a capped list of snippets." + ), + json_schema={ + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "document_ids": {"type": "array", "items": {"type": "string"}}, + "chunk_types": {"type": "array", "items": {"type": "string"}}, + "is_regex": {"type": "boolean", "default": False}, + "context_chars": {"type": "integer", "default": _DEFAULT_CONTEXT_CHARS}, + "max_results": {"type": "integer", "default": _DEFAULT_MAX_RESULTS}, + }, + "required": ["pattern"], + }, +) +async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + pattern = str(args.get("pattern") or "").strip() + if not pattern: + return ToolResult(text="", error="grep requires pattern") + is_regex = bool(args.get("is_regex", False)) + context_chars = int(args.get("context_chars") or _DEFAULT_CONTEXT_CHARS) + requested_max_results = int(args.get("max_results") or _DEFAULT_MAX_RESULTS) + max_results = capped_limit(requested_max_results, ctx.budget) + document_ids = [ + str(d).strip() for d in (args.get("document_ids") or []) if str(d).strip() + ] + chunk_types = { + str(t).strip().lower() for t in (args.get("chunk_types") or []) if str(t).strip() + } + + if is_regex: + try: + compiled = re.compile(pattern, flags=re.IGNORECASE) + except re.error as exc: + return ToolResult(text="", error=f"invalid regex: {exc}") + else: + compiled = re.compile(re.escape(pattern), flags=re.IGNORECASE) + + filters = _build_scope_filters( + user_id=ctx.user_id, + namespace=ctx.namespace, + document_ids=document_ids, + chunk_types=chunk_types, + ) + content_filter = ( + DocumentChunk.content.op("~*")(pattern) + if is_regex + else DocumentChunk.content.ilike(f"%{pattern}%") + ) + + count_stmt = ( + select(func.count(DocumentChunk.id)) + .select_from(DocumentChunk) + .join(Document, Document.document_id == DocumentChunk.document_id) + .where(*filters, content_filter) + ) + total_matches = int((await ctx.db.execute(count_stmt)).scalar_one()) + + rows_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.document_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentSection.section_path, + Document.source_file_name, + ) + .select_from(DocumentChunk) + .join(Document, Document.document_id == DocumentChunk.document_id) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .where(*filters, content_filter) + .order_by(DocumentChunk.document_id, DocumentChunk.sort_order) + .limit(max_results) + ) + rows = (await ctx.db.execute(rows_stmt)).all() + + results: list[dict[str, Any]] = [] + for chunk_id, document_id, chunk_type, content, section_path, source_file_name in rows: + text = str(content or "") + match = compiled.search(text) + snippet = build_snippet( + text, match.span() if match else None, hit_context=context_chars + ) + results.append( + { + "document_id": document_id, + "source_file_name": source_file_name, + "chunk_id": chunk_id, + "chunk_type": chunk_type, + "section_path": section_path, + "snippet": snippet, + } + ) + + lines = [f"total_matches={total_matches} returned={len(results)}"] + if requested_max_results > max_results: + lines.append(f"note: capped to budget.max_items={ctx.budget.max_items}") + for r in results: + lines.append(f"- {r['source_file_name']} / {r['section_path']}: {r['snippet']!r}") + + return ToolResult( + text="\n".join(lines), + payload={"total_matches": total_matches, "results": results}, + refs=[ + {"document_id": r["document_id"], "chunk_id": r["chunk_id"]} for r in results + ], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py new file mode 100644 index 00000000..d0ddffe1 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py @@ -0,0 +1,81 @@ +"""``corpus.list_documents`` — namespace-level document overview. + +Joins ``documents`` with the document-level ``graph_nodes`` row (§4 of +``CORPUS_SCHEMA.md``) to surface per-document keywords/summary/type-mix +without reading any chunk content. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select + +from shared.models.database.document import Document, GraphNode +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + register_tool, +) + + +@register_tool( + name="corpus.list_documents", + description=( + "List every active document in the namespace with its parse_track " + "and, when available, document-level graph metadata (top_keywords, " + "top_summary, chunk type mix). Use this to start cold: which " + "documents exist and what are they about, before picking one for " + "outline/node_filter/recall/read." + ), + json_schema={ + "type": "object", + "properties": {}, + "required": [], + }, +) +async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + stmt = ( + select(Document, GraphNode.properties) + .outerjoin( + GraphNode, + (GraphNode.owner_document_id == Document.document_id) + & (GraphNode.node_kind == "document"), + ) + .where(Document.user_id == ctx.user_id) + .where(Document.namespace == ctx.namespace) + .where(Document.status == "active") + .order_by(Document.source_file_name) + ) + rows = (await ctx.db.execute(stmt)).all() + + documents: list[dict[str, Any]] = [] + lines: list[str] = [] + for document, properties in rows: + props = properties if isinstance(properties, dict) else {} + entry = { + "document_id": document.document_id, + "source_file_name": document.source_file_name, + "parse_track": document.parse_track, + "top_keywords": props.get("top_keywords") or [], + "top_summary": props.get("top_summary") or "", + "types": props.get("types") or {}, + "chunks_count": props.get("chunks_count"), + } + documents.append(entry) + summary_line = ( + f"- {entry['source_file_name']} ({entry['document_id']}, " + f"track={entry['parse_track']}, chunks={entry['chunks_count']})" + ) + if entry["top_summary"]: + summary_line += f"\n summary: {entry['top_summary']}" + if entry["top_keywords"]: + summary_line += f"\n keywords: {', '.join(entry['top_keywords'])}" + lines.append(summary_line) + + text = f"documents={len(documents)}\n" + "\n".join(lines) if documents else "documents=0" + return ToolResult( + text=text, + payload={"documents": documents}, + refs=[{"document_id": doc["document_id"]} for doc in documents], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py new file mode 100644 index 00000000..2da21d68 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py @@ -0,0 +1,126 @@ +"""``corpus.neighbors`` — document-level ``related`` graph edges. + +Document-to-document only (§4 of ``CORPUS_SCHEMA.md``): no section- or +entity-level graph nodes exist yet. Edges are undirected and were written by +``DocumentGraphService.publish_document_graph`` with ``shared_entities`` (typed +entity overlap, preferred) or ``shared_keywords`` (TF-IDF fallback). +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import or_, select + +from shared.models.database.document import Document, GraphEdge, GraphNode +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + register_tool, +) + + +@register_tool( + name="corpus.neighbors", + description=( + "Return documents related to the given document via the persisted " + "document graph (typed-entity overlap, falling back to TF-IDF " + "keyword overlap), along with the shared terms that justify each " + "edge. Document-level only — there is no section- or entity-level " + "graph yet." + ), + json_schema={ + "type": "object", + "properties": {"document_id": {"type": "string"}}, + "required": ["document_id"], + }, +) +async def neighbors(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + document_id = str(args.get("document_id") or "").strip() + if not document_id: + return ToolResult(text="", error="neighbors requires document_id") + + document = ( + await ctx.db.execute( + select(Document) + .where(Document.document_id == document_id) + .where(Document.user_id == ctx.user_id) + .where(Document.namespace == ctx.namespace) + .where(Document.status == "active") + ) + ).scalar_one_or_none() + if document is None: + return ToolResult(text="", error=f"unknown document_id: {document_id}") + + node_id = f"doc:{document_id}" + edges = ( + ( + await ctx.db.execute( + select(GraphEdge) + .where(GraphEdge.user_id == ctx.user_id) + .where(GraphEdge.namespace == ctx.namespace) + .where(GraphEdge.edge_kind == "related") + .where( + or_( + GraphEdge.source_node_id == node_id, + GraphEdge.target_node_id == node_id, + ) + ) + ) + ) + .scalars() + .all() + ) + if not edges: + return ToolResult(text="neighbors=0", payload={"neighbors": []}) + + peer_node_ids = { + edge.target_node_id if edge.source_node_id == node_id else edge.source_node_id + for edge in edges + } + peer_nodes = ( + ( + await ctx.db.execute( + select(GraphNode).where(GraphNode.node_id.in_(peer_node_ids)) + ) + ) + .scalars() + .all() + ) + peer_by_id = {n.node_id: n for n in peer_nodes} + + neighbor_list: list[dict[str, Any]] = [] + for edge in sorted(edges, key=lambda e: -(e.weight or 0.0)): + peer_node_id = ( + edge.target_node_id if edge.source_node_id == node_id else edge.source_node_id + ) + peer = peer_by_id.get(peer_node_id) + if peer is None: + continue + props = edge.properties or {} + peer_props = peer.properties or {} + neighbor_list.append( + { + "document_id": peer.owner_document_id, + "source_file_name": peer_props.get("source_file_name"), + "weight": edge.weight, + "edge_basis": props.get("edge_basis"), + "shared_entities": props.get("shared_entities"), + "shared_keywords": props.get("shared_keywords"), + "connection_count": props.get("connection_count"), + } + ) + + lines = [f"neighbors={len(neighbor_list)}"] + for n in neighbor_list: + shared = n["shared_entities"] or n["shared_keywords"] or [] + lines.append( + f"- {n['source_file_name']} ({n['document_id']}) weight={n['weight']} " + f"basis={n['edge_basis']} shared={shared}" + ) + + return ToolResult( + text="\n".join(lines), + payload={"neighbors": neighbor_list}, + refs=[{"document_id": n["document_id"]} for n in neighbor_list], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py new file mode 100644 index 00000000..6c0269fa --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py @@ -0,0 +1,199 @@ +"""``corpus.node_filter`` — deterministic FOR-ALL/EXISTS/ANY/NOT predicate over sections. + +Reuses the exact predicate compile/match semantics from +``scoring.node_filter_predicates`` (path/summary substring|regex, fields AND together, +terms OR together) — see that module's docstring — but walks +``document_sections`` rows for the requested documents' current revision +instead of the in-memory map-nav tree. No top-K: returns the full matched set +and its count, per ``CORPUS_SCHEMA.md`` §6. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + register_tool, +) +from shared.services.retrieval.scoring.node_filter_predicates import ( + FieldPredicate, + _compile_predicates, + _node_matches, + field_predicate, +) + + +@register_tool( + name="corpus.node_filter", + description=( + "Deterministic FOR-ALL/EXISTS/ANY/NOT filter over section titles " + "(section_path) and summaries — not body text (use corpus.grep for " + "that). Predicates AND together across fields; terms within one " + "field's 'terms' list OR together. Returns the complete matched set " + "and its count, never a truncated top-K." + ), + json_schema={ + "type": "object", + "properties": { + "document_ids": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + }, + "predicates": { + "type": "array", + "items": { + "type": "object", + "properties": { + "field": {"type": "string", "enum": ["path", "summary"]}, + "terms": {"type": "array", "items": {"type": "string"}}, + "match": { + "type": "string", + "enum": ["substring", "regex"], + "default": "substring", + }, + }, + "required": ["field", "terms"], + }, + "minItems": 1, + }, + "chunk_types": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Narrow matched sections to those owning a body chunk of " + "one of these chunk_type values (e.g. ['page'] to filter " + "to page-track leaves only). Omit for no narrowing." + ), + }, + }, + "required": ["document_ids", "predicates"], + }, +) +async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + document_ids = [ + str(did).strip() for did in (args.get("document_ids") or []) if str(did).strip() + ] + if not document_ids: + return ToolResult(text="", error="node_filter requires document_ids") + raw_predicates = args.get("predicates") or [] + if not raw_predicates: + return ToolResult(text="", error="node_filter requires predicates") + + predicates: list[FieldPredicate] = [] + for raw in raw_predicates: + try: + predicates.append( + field_predicate( + raw.get("field"), + raw.get("terms") or [], + raw.get("match", "substring"), + ) + ) + except ValueError as exc: + return ToolResult(text="", error=str(exc)) + + compiled, failed = _compile_predicates(predicates) + if failed: + return ToolResult( + text=f"failed_predicates={failed}", + payload={"cardinality": 0, "matched_sections": [], "failed_predicates": failed}, + error="one or more predicates failed to compile", + ) + + documents = ( + ( + await ctx.db.execute( + select(Document) + .where(Document.document_id.in_(document_ids)) + .where(Document.user_id == ctx.user_id) + .where(Document.namespace == ctx.namespace) + .where(Document.status == "active") + ) + ) + .scalars() + .all() + ) + revision_by_doc = { + d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + } + if not revision_by_doc: + return ToolResult(text="", error="no active documents found for document_ids") + + revision_pairs = list(revision_by_doc.items()) + sections = ( + ( + await ctx.db.execute( + select(DocumentSection).where( + DocumentSection.document_id.in_([d for d, _ in revision_pairs]) + ) + ) + ) + .scalars() + .all() + ) + sections = [ + s for s in sections if revision_by_doc.get(s.document_id) == s.job_result_id + ] + + chunk_types = { + str(t).strip().lower() for t in (args.get("chunk_types") or []) if str(t).strip() + } + if chunk_types: + chunk_rows = await ctx.db.execute( + select(DocumentChunk.section_id, DocumentChunk.chunk_type).where( + DocumentChunk.document_id.in_([d for d, _ in revision_pairs]), + DocumentChunk.job_result_id.in_([r for _, r in revision_pairs]), + ) + ) + allowed_section_ids = { + str(section_id) + for section_id, chunk_type in chunk_rows.all() + if section_id and str(chunk_type or "").strip().lower() in chunk_types + } + sections = [s for s in sections if s.section_id in allowed_section_ids] + + matched_sections: list[dict[str, Any]] = [] + matched_doc_ids: list[str] = [] + seen_docs: set[str] = set() + for section in sorted(sections, key=lambda s: (s.document_id, s.sort_order)): + values = {"path": section.section_path, "summary": section.summary or ""} + if not _node_matches(values, compiled): + continue + matched_sections.append( + { + "document_id": section.document_id, + "section_id": section.section_id, + "section_path": section.section_path, + "summary": section.summary or "", + } + ) + if section.document_id not in seen_docs: + seen_docs.add(section.document_id) + matched_doc_ids.append(section.document_id) + + header = f"hits={len(matched_sections)}" + lines = [header] + for entry in matched_sections: + block = [entry["section_path"]] + if entry["summary"]: + block.append(f" summary: {entry['summary']}") + lines.append("\n".join(block)) + + return ToolResult( + text="\n".join(lines), + payload={ + "cardinality": len(matched_sections), + "matched_sections": matched_sections, + "matched_document_ids": matched_doc_ids, + }, + refs=[ + {"document_id": entry["document_id"], "section_path": entry["section_path"]} + for entry in matched_sections + ], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py new file mode 100644 index 00000000..db155d77 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py @@ -0,0 +1,155 @@ +"""``corpus.outline`` — titles + summaries, no body text, no folding. + +Reads ``document_sections`` (+ a ``document_chunks`` count aggregate) for one +document's current revision, optionally scoped to a ``section_path`` prefix +and depth-limited by the caller's own argument — never by a token budget +(see ``CORPUS_SCHEMA.md`` §6). This queries the live tables directly rather +than the compressed ``RetrievalNamespaceMapSnapshot``/serving-manifest blob: +that snapshot is namespace-wide and decoding it to read one document's +subtree would cost more than this document-scoped, index-backed query. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, select + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + register_tool, +) +from shared.services.retrieval.search.lexical_text import normalize_section_path + + +@register_tool( + name="corpus.outline", + description=( + "Return the section outline (title + summary + chunk_count) for one " + "document, or the subtree under a section_path prefix, without any " + "body text. Use for overviews, tables of contents, or picking where " + "to look before calling read. Depth is limited by the 'depth' " + "argument only, never truncated by a token budget." + ), + json_schema={ + "type": "object", + "properties": { + "document_id": {"type": "string"}, + "path_prefix": { + "type": "string", + "description": ( + "DB section_path (' / '-joined, e.g. 'Chapter 1 / Section " + "1.1'). Omit or use 'Root' for the whole document." + ), + }, + "depth": { + "type": "integer", + "description": ( + "Max levels below path_prefix (or below the document root " + "when path_prefix is omitted) to include. Omit for the " + "full subtree." + ), + }, + }, + "required": ["document_id"], + }, +) +async def outline(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + document_id = str(args.get("document_id") or "").strip() + if not document_id: + return ToolResult(text="", error="outline requires document_id") + depth_raw = args.get("depth") + depth = int(depth_raw) if depth_raw is not None else None + if depth is not None and depth < 0: + return ToolResult(text="", error="depth must be >= 0") + raw_prefix = str(args.get("path_prefix") or "").strip() + prefix = normalize_section_path(raw_prefix) if raw_prefix else "Root" + + document = ( + await ctx.db.execute( + select(Document) + .where(Document.document_id == document_id) + .where(Document.user_id == ctx.user_id) + .where(Document.namespace == ctx.namespace) + .where(Document.status == "active") + ) + ).scalar_one_or_none() + if document is None or not document.current_job_result_id: + return ToolResult(text="", error=f"unknown document_id: {document_id}") + job_result_id = document.current_job_result_id + + section_stmt = ( + select(DocumentSection) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .order_by(DocumentSection.sort_order, DocumentSection.section_id) + ) + all_sections = list((await ctx.db.execute(section_stmt)).scalars().all()) + + base_level: int | None = None + if prefix == "Root": + base_level = 0 + scoped = all_sections + else: + prefix_section = next( + (s for s in all_sections if s.section_path == prefix), None + ) + if prefix_section is None: + return ToolResult( + text="", error=f"unknown path_prefix for {document_id}: {prefix}" + ) + base_level = prefix_section.section_level + scoped = [ + s + for s in all_sections + if s.section_path == prefix or s.section_path.startswith(f"{prefix} / ") + ] + + if depth is not None: + scoped = [s for s in scoped if (s.section_level - base_level) <= depth] + + section_ids = [s.section_id for s in scoped] + chunk_counts: dict[str, int] = {} + if section_ids: + count_rows = await ctx.db.execute( + select(DocumentChunk.section_id, func.count(DocumentChunk.id)) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(section_ids)) + .group_by(DocumentChunk.section_id) + ) + chunk_counts = {str(sid): int(count) for sid, count in count_rows.all()} + + nodes: list[dict[str, Any]] = [] + lines: list[str] = [] + for section in scoped: + relative_depth = section.section_level - base_level + node = { + "section_id": section.section_id, + "section_path": section.section_path, + "section_title": section.section_title, + "section_level": section.section_level, + "relative_depth": relative_depth, + "summary": section.summary or "", + "chunk_count": chunk_counts.get(section.section_id, 0), + } + nodes.append(node) + indent = " " * max(relative_depth, 0) + line = f"{indent}- {node['section_title'] or node['section_path']} (chunks={node['chunk_count']})" + if node["summary"]: + line += f"\n{indent} summary: {node['summary']}" + lines.append(line) + + text = f"document={document.source_file_name} sections={len(nodes)}\n" + "\n".join( + lines + ) + return ToolResult( + text=text, + payload={"document_id": document_id, "sections": nodes}, + refs=[ + {"document_id": document_id, "section_path": node["section_path"]} + for node in nodes + ], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py new file mode 100644 index 00000000..f5f99117 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py @@ -0,0 +1,399 @@ +"""``corpus.read`` — full body content for already-located sections/chunks. + +Unlike ``hydration.result_assembly.assemble_retrieval_results`` (which +down-weights ``page`` chunks to their summary — see that module's +``_page_summary``, a deliberate trade-off for the retrieval-answer surface), +``read`` returns the page chunk's full body content, with ``[SAME-AS +p]`` markers resolved to the owner section's text (§2 of +``CORPUS_SCHEMA.md``) rather than stripped or summarized. ``connect_to`` +assets are still inlined via the same placeholder mechanism as retrieval, and +``page_assets``/asset ``file_path`` are converted to URLs via the existing +``enrich_rows_with_retrieval_asset_url``. + +SAME-AS resolution is single-level: the owner chunk's full content is +embedded as-is. If that owner chunk itself still contains an unrelated +SAME-AS marker (a different leaf's page), it is not recursively resolved in +this pass — a disclosed scope limit, not a silent gap (the raw marker stays +visible in the embedded text). + +``section_path`` refs are resolved via ``agent_tools.section_path_lookup``: +exact match first, then a unique segment-bound suffix match when the agent +omits ancestor segments; ambiguous suffix matches return an error listing +candidate full paths instead of picking one silently. +""" + +from __future__ import annotations + +import re +from typing import Any + +from sqlalchemy import select + +from shared.models.database.document import ( + Document, + DocumentChunk, + DocumentSection, +) +from shared.models.database.job_result import JobResult +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + register_tool, +) +from shared.services.retrieval.hydration.assets import ( + enrich_rows_with_retrieval_asset_url, +) +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.result_assembly import ( + _compose_table_content, + _compose_text_content, + _image_display_content, +) +from shared.services.retrieval.hydration.row_utils import normalize_chunk_type +from shared.services.retrieval.agent_tools.section_path_lookup import ( + resolve_section_path_anchor, + section_path_anchor_filter, + section_path_subtree_filter, +) +from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path + +_SAME_AS_MARKER_RE = re.compile(r"\[SAME-AS (.+?) p(\d+)\]") +_BODY_CHUNK_TYPES = ("text", "page") + +# ``_compose_text_content`` doesn't branch on chunk_type — it just inlines +# connect_to placeholders — so the ``page`` branch below reuses it directly +# instead of carrying a near-identical copy. The behavioral difference from +# retrieval's own page handling (never downgrading to a summary — see the +# module docstring) comes entirely from *not* calling ``_page_summary`` +# first, which this module never did. + + +async def _resolve_same_as_markers( + db: Any, + rows: list[dict[str, Any]], + *, + revision_by_doc: dict[str, str], + source_file_name_by_doc: dict[str, str], +) -> None: + """Mutate ``page`` rows in place, replacing SAME-AS markers with owner text.""" + matches_by_index: dict[int, list[tuple[str, str]]] = {} + needed: set[tuple[str, str]] = set() + for index, row in enumerate(rows): + if normalize_chunk_type(row.get("chunk_type")) != "page": + continue + content = str(row.get("content") or "") + found = list(_SAME_AS_MARKER_RE.finditer(content)) + if not found: + continue + document_id = str(row.get("document_id") or "") + source_file_name = source_file_name_by_doc.get(document_id) + row_matches: list[tuple[str, str]] = [] + for match in found: + owner_db_path = section_path_from_chunk_path( + match.group(1), source_file_name=source_file_name + ) + needed.add((document_id, owner_db_path)) + row_matches.append((match.group(0), owner_db_path)) + matches_by_index[index] = row_matches + + if not needed: + return + + owner_content: dict[tuple[str, str], str] = {} + for document_id, owner_path in needed: + job_result_id = revision_by_doc.get(document_id) + if not job_result_id: + continue + result = await db.execute( + select(DocumentChunk.content) + .select_from(DocumentChunk) + .join(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentSection.section_path == owner_path) + .where(DocumentChunk.chunk_type == "page") + ) + content_row = result.first() + owner_content[(document_id, owner_path)] = ( + str(content_row[0]) if content_row and content_row[0] else "" + ) + + for index, row_matches in matches_by_index.items(): + row = rows[index] + content = str(row.get("content") or "") + document_id = str(row.get("document_id") or "") + for marker_text, owner_path in row_matches: + resolved = owner_content.get((document_id, owner_path), "") + if resolved: + replacement = f"(SAME-AS {owner_path} resolved)\n{resolved}" + else: + replacement = f"(SAME-AS {owner_path} — page not found)" + content = content.replace(marker_text, replacement, 1) + row["content"] = content + + +@register_tool( + name="corpus.read", + description=( + "Read full body content for already-located sections or chunks. " + "Resolves page-track SAME-AS pointers to the owner section's text, " + "inlines connect_to assets, and converts asset/page_assets " + "references to URLs. Use after outline/node_filter/recall/grep " + "have located where to look." + ), + json_schema={ + "type": "object", + "properties": { + "refs": { + "type": "array", + "items": { + "type": "object", + "properties": { + "document_id": {"type": "string"}, + "section_path": {"type": "string"}, + "chunk_id": {"type": "string"}, + }, + "required": ["document_id"], + }, + "minItems": 1, + "description": "Each ref needs document_id and either section_path or chunk_id.", + }, + "mode": { + "type": "string", + "enum": ["self", "descendants"], + "default": "self", + "description": ( + "'descendants' also reads every section under a " + "section_path ref; ignored for chunk_id refs." + ), + }, + "include_assets": {"type": "boolean", "default": True}, + "resolve_same_as": {"type": "boolean", "default": True}, + }, + "required": ["refs"], + }, +) +async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + refs = args.get("refs") or [] + if not refs: + return ToolResult(text="", error="read requires refs") + mode = str(args.get("mode") or "self").strip().lower() + if mode not in ("self", "descendants"): + return ToolResult(text="", error=f"unsupported mode: {mode}") + include_assets = bool(args.get("include_assets", True)) + resolve_same_as_flag = bool(args.get("resolve_same_as", True)) + + document_ids = { + str(ref.get("document_id") or "").strip() for ref in refs if ref.get("document_id") + } + documents = ( + ( + await ctx.db.execute( + select(Document) + .where(Document.document_id.in_(document_ids)) + .where(Document.user_id == ctx.user_id) + .where(Document.namespace == ctx.namespace) + .where(Document.status == "active") + ) + ) + .scalars() + .all() + ) + revision_by_doc = { + d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + } + source_file_name_by_doc = {d.document_id: d.source_file_name or "" for d in documents} + job_result_ids = sorted(set(revision_by_doc.values())) + job_id_by_revision: dict[str, str] = {} + if job_result_ids: + job_rows = await ctx.db.execute( + select(JobResult.id, JobResult.job_id).where(JobResult.id.in_(job_result_ids)) + ) + job_id_by_revision = {str(rid): str(jid) for rid, jid in job_rows.all() if rid and jid} + + base_rows: list[dict[str, Any]] = [] + errors: list[str] = [] + for ref in refs: + document_id = str(ref.get("document_id") or "").strip() + job_result_id = revision_by_doc.get(document_id) + if not job_result_id: + errors.append(f"unknown document_id: {document_id}") + continue + chunk_id = str(ref.get("chunk_id") or "").strip() + section_path = str(ref.get("section_path") or "").strip() + source_file_name = source_file_name_by_doc.get(document_id, "") + job_id = job_id_by_revision.get(job_result_id) + + if chunk_id: + row = ( + await ctx.db.execute( + select(DocumentChunk, DocumentSection.section_path) + .select_from(DocumentChunk) + .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_id == chunk_id) + ) + ).first() + if row is None: + errors.append(f"unknown chunk_id: {chunk_id} in {document_id}") + continue + chunk, resolved_section_path = row + base_rows.append( + { + "document_id": document_id, + "job_result_id": job_result_id, + "job_id": job_id, + "source_file_name": source_file_name, + "chunk_id": chunk.chunk_id, + "section_id": chunk.section_id, + "section_path": resolved_section_path, + "chunk_type": chunk.chunk_type, + "content": chunk.content, + "chunk_metadata": chunk.chunk_metadata or {}, + "file_path": chunk.file_path, + } + ) + continue + + if not section_path: + errors.append(f"ref for {document_id} needs section_path or chunk_id") + continue + + resolved_path, path_error = await resolve_section_path_anchor( + ctx.db, + document_id=document_id, + job_result_id=job_result_id, + section_path=section_path, + ) + if path_error or not resolved_path: + errors.append(path_error or f"unknown section_path for {document_id}") + continue + + path_filter = ( + section_path_subtree_filter(resolved_path) + if mode == "descendants" + else section_path_anchor_filter(resolved_path) + ) + section_rows = ( + ( + await ctx.db.execute( + select(DocumentSection) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .where(path_filter) + .order_by(DocumentSection.sort_order) + ) + ) + .scalars() + .all() + ) + section_ids = [s.section_id for s in section_rows] + chunk_rows = ( + await ctx.db.execute( + select(DocumentChunk).where( + DocumentChunk.document_id == document_id, + DocumentChunk.job_result_id == job_result_id, + DocumentChunk.section_id.in_(section_ids), + # Body chunks only (text/page). image/table chunks share a + # section_id with whichever section happens to store them + # in the DB (always Root — CORPUS_SCHEMA.md §3), which is + # not the same as "belonging" to that section; their real + # association is connect_to on the body chunk, resolved + # below via hydrate_connected_target_rows. Without this + # filter, reading Root would return every still-unmounted + # asset in the document as spurious top-level entries. + DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES), + ) + ) + ).scalars().all() + section_path_by_id = {s.section_id: s.section_path for s in section_rows} + for chunk in chunk_rows: + base_rows.append( + { + "document_id": document_id, + "job_result_id": job_result_id, + "job_id": job_id, + "source_file_name": source_file_name, + "chunk_id": chunk.chunk_id, + "section_id": chunk.section_id, + "section_path": ( + section_path_by_id.get(chunk.section_id) + if chunk.section_id + else None + ), + "chunk_type": chunk.chunk_type, + "content": chunk.content, + "chunk_metadata": chunk.chunk_metadata or {}, + "file_path": chunk.file_path, + } + ) + + if not base_rows: + return ToolResult( + text="", + error="no chunks resolved for given refs" + (f" ({'; '.join(errors)})" if errors else ""), + ) + + if resolve_same_as_flag: + await _resolve_same_as_markers( + ctx.db, + base_rows, + revision_by_doc=revision_by_doc, + source_file_name_by_doc=source_file_name_by_doc, + ) + + connected_rows: list[dict[str, Any]] = [] + if include_assets: + connected_rows = await hydrate_connected_target_rows( + db=ctx.db, + rows=base_rows, + exclude_document_ids=[], + exclude_sections=[], + ) + rows_by_chunk_id = { + str(row.get("chunk_id") or ""): row + for row in [*base_rows, *connected_rows] + if row.get("chunk_id") + } + + assembled: list[dict[str, Any]] = [] + for row in base_rows: + chunk_type = normalize_chunk_type(row.get("chunk_type")) + composed = dict(row) + if chunk_type == "text": + composed["content"] = _compose_text_content(row, rows_by_chunk_id) if include_assets else row.get("content") + elif chunk_type == "page": + composed["content"] = _compose_text_content(row, rows_by_chunk_id) if include_assets else row.get("content") + elif chunk_type == "table": + composed["content"] = _compose_table_content(row, rows_by_chunk_id) + elif chunk_type == "image": + composed["content"] = _image_display_content(row) + assembled.append(composed) + + if include_assets: + assembled = await enrich_rows_with_retrieval_asset_url( + assembled, log_context="agent_tools.read" + ) + + lines = [] + if errors: + lines.append(f"errors: {'; '.join(errors)}") + for row in assembled: + lines.append( + f"### {row.get('source_file_name')} / {row.get('section_path')} " + f"[{row.get('chunk_type')}]" + ) + lines.append(str(row.get("content") or "")) + + return ToolResult( + text="\n".join(lines), + payload={"chunks": assembled, "errors": errors}, + refs=[ + {"document_id": row["document_id"], "chunk_id": row["chunk_id"]} + for row in assembled + ], + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py new file mode 100644 index 00000000..6e906cd6 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py @@ -0,0 +1,270 @@ +"""``corpus.recall`` — fuzzy ranked candidate search. + +Two real channels today, fused by the existing RRF utility +(``search.scoring.merge_channels_rrf``, same ``RRF_K`` as every other RRF +fusion in retrieval): + +- ``path_content``: reuses ``search.map_unit_discovery.map_unit_discovery`` + (persisted map-unit BM25 over path+content, already RRF-fused internally). +- ``term``: a fresh substring channel over + ``document_map_units.term_search_text_lower`` — this column is persisted at + index time and is ranked here as an independent substring channel. + +``vector`` is accepted in ``channels`` but rejected as reserved/not +implemented (``CORPUS_SCHEMA.md`` §5) — it is not silently ignored. + +Fusing an already-doubly-fused channel (path_content) with a fresh single +channel (term) at equal RRF weight is a necessary, disclosed design choice: +there is no persisted precedent for a different weight ratio between them +(the old 3-channel weights of path=1.0/content=2.0/term=1.5 no longer exist +in code — only path=1.0/content=2.0 survive in ``scoring.knowhere_hybrid``). + +The term channel's snippet and the rendered ``text`` preview both go through +the shared ``agent_tools.snippet.build_snippet`` (head + first-match window + +tail, ``...``-joined, overlap-merged) — the same mechanism ``corpus.grep`` +uses, so window-slicing constants live in one place. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select, text +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentChunk +from shared.services.retrieval.agent_tools.registry import ( + ToolContext, + ToolResult, + capped_limit, + register_tool, +) +from shared.services.retrieval.agent_tools.snippet import build_snippet +from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery +from shared.services.retrieval.search.scoring import merge_channels_rrf + +_SUPPORTED_CHANNELS = {"path_content", "term"} +_RESERVED_CHANNELS = {"vector"} +_DEFAULT_TOP_K = 10 + +_TERM_CHANNEL_SQL = """ +SELECT dmu.document_id, dmu.job_result_id, dmu.section_id, ds.section_path, + d.source_file_name, dmu.term_search_text_lower +FROM document_map_units dmu +JOIN documents d + ON d.document_id = dmu.document_id + AND d.current_job_result_id = dmu.job_result_id +JOIN document_sections ds ON ds.section_id = dmu.section_id +WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + {doc_clause} + AND dmu.term_search_text_lower LIKE :like_pattern +ORDER BY POSITION(:needle IN dmu.term_search_text_lower) ASC +LIMIT :limit +""" + + +async def _excluded_document_ids( + db: AsyncSession, *, user_id: str, namespace: str, document_ids: list[str] +) -> list[str]: + if not document_ids: + return [] + rows = await db.execute( + select(Document.document_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .where(Document.document_id.notin_(document_ids)) + ) + return [str(r[0]) for r in rows.all()] + + +async def _term_channel_rows( + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + document_ids: list[str], + chunk_types: set[str] | None, + top_k: int, +) -> list[dict[str, Any]]: + needle = query.strip().lower() + if not needle: + return [] + params: dict[str, Any] = { + "user_id": user_id, + "namespace": namespace, + "like_pattern": f"%{needle}%", + "needle": needle, + "limit": top_k, + } + doc_clause = "" + if document_ids: + doc_clause = "AND d.document_id = ANY(:doc_ids)" + params["doc_ids"] = document_ids + statement = text(_TERM_CHANNEL_SQL.format(doc_clause=doc_clause)) + unit_rows = [dict(row._mapping) for row in (await db.execute(statement, params)).all()] + if not unit_rows: + return [] + + keys = [ + (row["document_id"], row["job_result_id"], row["section_id"]) for row in unit_rows + ] + chunk_result = await db.execute( + select(DocumentChunk).where( + DocumentChunk.document_id.in_({k[0] for k in keys}), + DocumentChunk.job_result_id.in_({k[1] for k in keys}), + DocumentChunk.section_id.in_({k[2] for k in keys}), + ) + ) + chunk_by_key = { + (c.document_id, c.job_result_id, c.section_id): c + for c in chunk_result.scalars().all() + } + + results: list[dict[str, Any]] = [] + for unit_row in unit_rows: + key = (unit_row["document_id"], unit_row["job_result_id"], unit_row["section_id"]) + chunk = chunk_by_key.get(key) + if chunk is None: + continue + if chunk_types and chunk.chunk_type not in chunk_types: + continue + haystack = unit_row["term_search_text_lower"] + needle_pos = haystack.find(needle) + hit = (needle_pos, needle_pos + len(needle)) if needle_pos >= 0 else None + results.append( + { + "chunk_id": chunk.chunk_id, + "document_id": chunk.document_id, + "section_id": chunk.section_id, + "section_path": unit_row["section_path"], + "source_file_name": unit_row["source_file_name"], + "chunk_type": chunk.chunk_type, + "snippet": build_snippet(haystack, hit), + } + ) + return results + + +@register_tool( + name="corpus.recall", + description=( + "Fuzzy ranked candidate search for a question when you don't know " + "where the answer lives. Fuses a path+content BM25 channel with a " + "term substring channel via RRF. Returns candidates with path and " + "snippet, not full content — call corpus.read on the winners." + ), + json_schema={ + "type": "object", + "properties": { + "query": {"type": "string"}, + "document_ids": {"type": "array", "items": {"type": "string"}}, + "chunk_types": {"type": "array", "items": {"type": "string"}}, + "channels": { + "type": "array", + "items": { + "type": "string", + "enum": ["path_content", "term", "vector"], + }, + "default": ["path_content", "term"], + "description": "'vector' is reserved and not implemented yet.", + }, + "top_k": {"type": "integer", "default": _DEFAULT_TOP_K}, + }, + "required": ["query"], + }, +) +async def recall(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + query = str(args.get("query") or "").strip() + if not query: + return ToolResult(text="", error="recall requires query") + requested_top_k = int(args.get("top_k") or _DEFAULT_TOP_K) + top_k = capped_limit(requested_top_k, ctx.budget) + document_ids = [ + str(d).strip() for d in (args.get("document_ids") or []) if str(d).strip() + ] + chunk_types = { + str(t).strip().lower() for t in (args.get("chunk_types") or []) if str(t).strip() + } or None + requested_channels = set(args.get("channels") or list(_SUPPORTED_CHANNELS)) + reserved_requested = requested_channels & _RESERVED_CHANNELS + active_channels = requested_channels & _SUPPORTED_CHANNELS + unknown_channels = requested_channels - _SUPPORTED_CHANNELS - _RESERVED_CHANNELS + if unknown_channels: + return ToolResult(text="", error=f"unsupported channels: {sorted(unknown_channels)}") + if not active_channels: + return ToolResult( + text="", + error="no runnable channels requested (vector is reserved, not implemented)", + ) + + channel_rows: list[list[dict[str, Any]]] = [] + weights: list[float] = [] + + if "path_content" in active_channels: + exclude_document_ids = await _excluded_document_ids( + ctx.db, user_id=ctx.user_id, namespace=ctx.namespace, document_ids=document_ids + ) + discovery = await map_unit_discovery( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + query=query, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=[], + chunk_types=chunk_types, + ) + channel_rows.append(list(discovery.payload.get("fused_rows") or [])) + weights.append(1.0) + + if "term" in active_channels: + term_rows = await _term_channel_rows( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + query=query, + document_ids=document_ids, + chunk_types=chunk_types, + top_k=top_k, + ) + channel_rows.append(term_rows) + weights.append(1.0) + + fused = merge_channels_rrf(channel_rows, weights, top_k) + + lines = [f"candidates={len(fused)}"] + if len(fused) < 2: + # No tool name named here on purpose — this fires on *every* weak + # recall regardless of what a better next step happens to be for + # this corpus/query, so it nudges the agent to change approach + # without prescribing which other tool to reach for (that's already + # covered generically in CORPUS_SCHEMA.md §6's tool-selection table). + lines.append( + "note: few or no candidates for this phrasing — rephrasing the " + "query and calling recall again rarely surfaces more; a " + "different exploration approach is more likely to help than " + "repeating recall with synonyms." + ) + if reserved_requested: + lines.append(f"note: channels {sorted(reserved_requested)} are reserved, not run") + if requested_top_k > top_k: + lines.append(f"note: capped to budget.max_items={ctx.budget.max_items}") + for row in fused: + snippet = build_snippet(str(row.get("content") or row.get("snippet") or "")) + lines.append( + f"- {row.get('source_file_name')} / {row.get('section_path')} " + f"score={row.get('score')}: {snippet!r}" + ) + + return ToolResult( + text="\n".join(lines), + payload={"candidates": fused, "reserved_channels": sorted(reserved_requested)}, + refs=[ + {"document_id": row.get("document_id"), "chunk_id": row.get("chunk_id")} + for row in fused + ], + ) diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index c327eb46..07ea6b6a 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -79,6 +79,7 @@ def _cache_shape_digest( use_agentic: bool | None = None, llm_text_model: str | None = None, llm_vision_model: str | None = None, + harness: str | None = None, ) -> str: normalized_excludes = sorted(exclude_document_ids) normalized_sections = _normalize_exclude_sections(exclude_sections) @@ -96,6 +97,7 @@ def _cache_shape_digest( str(use_agentic), str(llm_text_model or ""), str(llm_vision_model or ""), + str(harness or ""), ] ) payload = f"{query}|{top_k}|{'|'.join(normalized_excludes)}|{'|'.join(normalized_sections)}|{extra}" diff --git a/packages/shared-python/shared/services/retrieval/execution/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py index 9e831bce..c6d3488d 100644 --- a/packages/shared-python/shared/services/retrieval/execution/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -8,6 +8,7 @@ from shared.models.schemas.llm_config import LLMConfig from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.execution.route_types import RetrievalRouteContext +from shared.services.retrieval.agent_explore.harness.resolve import resolve_harness_name from shared.services.retrieval.settings import ( INTERNAL_RECALL_K_MULTIPLIER, ) @@ -98,6 +99,9 @@ def build_cache_extra(self) -> dict[str, Any]: "use_agentic": self.use_agentic, "llm_text_model": text_model, "llm_vision_model": vision_model, + "harness": ( + resolve_harness_name() if self.use_agentic is not False else "classic" + ), } def resolve_allowed_chunk_types(self) -> set[str] | None: diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py index 656533c6..1ac29076 100644 --- a/packages/shared-python/shared/services/retrieval/execution/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py @@ -16,7 +16,7 @@ def to_public_source(row: dict[str, Any]) -> dict[str, Any]: async def enrich_referenced_chunks_with_asset_url(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: return await enrich_rows_with_retrieval_asset_url( refs, - log_context='mapnav referenced chunk', + log_context='referenced chunk', ) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 3f6a8270..b14a84e9 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -1,7 +1,5 @@ from __future__ import annotations -import asyncio -import resource import time from contextlib import AbstractAsyncContextManager @@ -25,10 +23,6 @@ count_scoped_chunks, load_all_scoped_chunks, ) -from shared.services.retrieval.execution.revision_pins import ( - capture_revision_pins, - is_revision_generation_stable, -) def open_fresh_database_context() -> AbstractAsyncContextManager[AsyncSession]: @@ -67,11 +61,10 @@ async def run_retrieval_route( if small_corpus_outcome is not None: return small_corpus_outcome - # Explicit False → classic 3-channel top-K. None/True → map-nav (default). + # Explicit False → classic map-unit BM25 top-K. None/True → agent_explore. if context.use_agentic is False: return await _run_classic_topk_route(context) - - return await _run_mapnav_route(context) + return await _run_agent_explore_route(context) async def _try_run_small_corpus_route( @@ -190,198 +183,114 @@ async def _run_classic_topk_route( ) -async def _run_mapnav_route( +async def _run_agent_explore_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: - """Default agentic path: PLANNER + HARVEST + CONTROL (checklist map-nav). + """Default agentic path: in-process ``corpus.*`` tool-calling loop. - LEGACY, PENDING REPLACEMENT: ``agent_explore`` will become the default - agentic route once it passes its evaluation gate; this route then stays - only as the ``RETRIEVAL_AGENTIC_ROUTER=mapnav`` fallback until Phase 5 - cleanup. Do not add new capabilities here — new agentic-retrieval work - belongs in ``shared/services/retrieval/agent_tools/`` and - ``shared/services/retrieval/agent_explore/``. + Which provider actually runs the tool-calling loop is the + ``AGENT_EXPLORE_HARNESS`` switch resolved by ``resolve_harness()``. """ - process_started = resource.getrusage(resource.RUSAGE_SELF) - from shared.services.retrieval import nav_llm_backend # noqa: F401 - from shared.services.retrieval.nav import run_nav_episode - from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace - from shared.services.retrieval.nav_bridge import build_referenced_chunks - from shared.services.retrieval.nav_config import ( - MAPNAV_MODEL, - build_nav_config, - nav_evidence_chars, + from shared.services.retrieval.agent_explore.bridge import build_decision_trace + from shared.services.retrieval.agent_explore.budget import EpisodeBudget + from shared.services.retrieval.agent_explore.harness import resolve_harness + from shared.services.retrieval.agent_explore.ref_resolution import ( + resolve_finish_refs, ) - from shared.services.retrieval.nav_snapshot import load_nav_snapshot - from shared.services.retrieval.trace import ( - TraceRecorder, - build_decision_trace, - episode_selected_doc_ids, - episode_selected_paths, - episode_token_count, - episode_workflow_plan, + from shared.services.retrieval.trace import TraceRecorder + + harness = resolve_harness() + episode_started = time.perf_counter() + episode = await harness.run_episode( + db_factory=open_fresh_database_context, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + budget=EpisodeBudget(), + ) + logger.info( + "retrieval agent_explore stage=episode seconds={:.3f} refs={} " + "steps={} tokens={} stop_reason={}".format( + time.perf_counter() - episode_started, + len(episode.refs), + len(episode.steps), + episode.tokens_used, + episode.stop_reason, + ) ) - snapshot_started = time.perf_counter() - snapshot_pins = context.revision_pins - snapshot = await load_nav_snapshot( + # episode.refs are document_id + section_path (what the agent actually + # sees in tool text); resolve_workflow_references requires chunk_id — + # see ref_resolution.py's module docstring for why this bridge exists. + chunk_refs = await resolve_finish_refs( context.db, user_id=context.user_id, namespace=context.namespace, + refs=episode.refs, + ) + resolved = await resolve_workflow_references( + db=context.db, + user_id=context.user_id, + namespace=context.namespace, + refs=chunk_refs, + revision_pins=context.revision_pins, + ) + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=resolved.rows, exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, - lazy=True, - revision_pins=snapshot_pins, - generation=(snapshot_pins.generation if snapshot_pins is not None else None), + allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, + ) + + decision_steps = build_decision_trace(episode.steps) + decision_trace = [step.to_dict() for step in decision_steps] + selected_doc_ids = list( + {row.get("document_id", "") for row in resolved.rows if row.get("document_id")} ) - if snapshot_pins is not None and not await is_revision_generation_stable( + + trace = TraceRecorder( context.db, user_id=context.user_id, namespace=context.namespace, - pins=snapshot_pins, - ): - snapshot.close() - snapshot_pins = await capture_revision_pins( - context.db, - user_id=context.user_id, - namespace=context.namespace, - ) - snapshot = await load_nav_snapshot( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - lazy=True, - revision_pins=snapshot_pins, - generation=(snapshot_pins.generation if snapshot_pins is not None else None), - ) - snapshot_seconds = time.perf_counter() - snapshot_started - logger.info( - "retrieval mapnav stage=snapshot_load seconds={:.3f} documents={} refs={} " - "conversation_id={}".format( - snapshot_seconds, - len(snapshot.document_ids), - len(snapshot.chunk_ref_index), - context.conversation_id or "", - ) + query=context.query, + top_k=context.top_k, + chunk_types=context.allowed_chunk_types, + policy_name="agent_explore_v1", ) - - # Small-corpus count / snapshot reads may leave a checkout; drop it before - # the sync LLM episode (same pattern as the retired workflow route). - await context.db.rollback() - - budget = nav_evidence_chars() - cfg = build_nav_config() - toolspace = ProviderToolSpace(snapshot.provider) - - episode_started = time.perf_counter() - try: - episode = await asyncio.to_thread( - run_nav_episode, - None, - context.query, - corpus_doc_ids=list(snapshot.document_ids), - budget_chars=budget, - compose_answer=False, - policy="llm", - config=cfg, - toolspace=toolspace, - ) - - refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot) - logger.info( - "retrieval mapnav stage=episode seconds={:.3f} refs={}".format( - time.perf_counter() - episode_started, - len(refs), - ) - ) - finally: - snapshot.close() - - hydration_started = time.perf_counter() - async with open_fresh_database_context() as final_db: - resolved = await resolve_workflow_references( - db=final_db, - user_id=context.user_id, - namespace=context.namespace, - refs=refs, - score_by_chunk_id=score_by_chunk_id or None, - revision_pins=snapshot.document_revisions, - ) - assembled_rows = await assemble_retrieval_results( - db=final_db, - rows=resolved.rows, - exclude_document_ids=context.exclude_document_ids, - exclude_sections=context.exclude_sections, - allowed_chunk_types=context.allowed_chunk_types, - revision_pins=snapshot.document_revisions, - ) - - decision_steps = build_decision_trace( - episode, - evidence_char_budget=budget, - n_refs=len(resolved.refs), - ) - decision_trace = [step.to_dict() for step in decision_steps] - selected_paths = episode_selected_paths(episode, resolved.refs) - selected_docs = episode_selected_doc_ids(resolved.refs) - tokens_used = episode_token_count(episode) - trace = TraceRecorder( - final_db, - user_id=context.user_id, - namespace=context.namespace, - query=context.query, - top_k=context.top_k, - chunk_types=context.allowed_chunk_types, - workflow_plan=episode_workflow_plan(episode), - policy_name="mapnav_checklist_v1", - ) - await trace.create_run() - for step in decision_steps: - trace.record_decision_trace_step(step) - await trace.complete( - assembled_rows, - "mapnav", - token_count=tokens_used, - model_name=MAPNAV_MODEL, - selected_paths=selected_paths, - selected_doc_ids=selected_docs, - ) - logger.info( - "retrieval mapnav stage=hydration seconds={:.3f} results={}".format( - time.perf_counter() - hydration_started, - len(assembled_rows), - ) + await trace.create_run() + for step in decision_steps: + trace.record_decision_trace_step(step) + if episode.stop_reason.startswith("budget_"): + trace.record_budget_stop(episode.stop_reason.removeprefix("budget_")) + await trace.complete( + assembled_rows, + "agent_explore", + token_count=episode.tokens_used, + model_name=episode.model_name, + selected_doc_ids=selected_doc_ids, ) - stop_reason = str(getattr(episode, "stop_reason", "") or "completed") - evidence_text = str(getattr(episode, "evidence_text", "") or "") + evidence_text = _render_rows_evidence(assembled_rows) response = { "namespace": context.namespace, "query": context.query, - "router_used": "mapnav", + "router_used": "agent_explore", "evidence_text": evidence_text, "answer_text": "", "referenced_chunks": resolved.refs, "results": assembled_rows, - "stop_reason": stop_reason, + "stop_reason": episode.stop_reason, "decision_trace": decision_trace, } - - completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" - process_finished = resource.getrusage(resource.RUSAGE_SELF) - logger.info( - "retrieval mapnav stage=process_resources cpu_seconds={:.3f} " - "process_max_rss_kb={}", - (process_finished.ru_utime + process_finished.ru_stime) - - (process_started.ru_utime + process_started.ru_stime), - int(process_finished.ru_maxrss), - ) return RetrievalRouteOutcome( response=response, hit_stats_results=resolved.refs, - completion_label="MAPNAV RETRIEVAL", + completion_label="AGENT EXPLORE RETRIEVAL", completion_count=len(resolved.refs), - completion_detail=completion_detail, + completion_detail=( + f"chunks | evidence={len(evidence_text)} chars | router=agent_explore" + ), ) + diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index ab7a0523..c58e5f8e 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -1,4 +1,4 @@ -"""Publication-time materialization of exact map-nav lexical units.""" +"""Publication-time materialization of persisted map-unit lexical units.""" from __future__ import annotations @@ -16,15 +16,15 @@ DocumentMapUnitToken, DocumentSection, ) -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION +from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, SectionRow, UnitRow, ) -from shared.services.retrieval.nav.nav_map_scores import build_score_units -from shared.services.retrieval.nav.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION -from shared.services.retrieval.nav.persisted_score_load import average_idf_from_unit_dfs +from shared.services.retrieval.scoring.persisted_score_load import average_idf_from_unit_dfs +from shared.services.retrieval.scoring.score_units import build_score_units from shared.services.retrieval.publication_models import DocumentPublicationScope __all__ = ["MAP_UNIT_INDEX_FORMAT_VERSION", "replace_document_map_units"] diff --git a/packages/shared-python/shared/services/retrieval/scoring/__init__.py b/packages/shared-python/shared/services/retrieval/scoring/__init__.py new file mode 100644 index 00000000..404eaf74 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoring/__init__.py @@ -0,0 +1 @@ +"""Shared retrieval scoring primitives used by publication and classic recall.""" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/scoring/hierarchy.py similarity index 75% rename from packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py rename to packages/shared-python/shared/services/retrieval/scoring/hierarchy.py index 577de014..f9dcfe0d 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/scoring/hierarchy.py @@ -21,7 +21,7 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import ( Any, Dict, @@ -41,7 +41,17 @@ ) if TYPE_CHECKING: - from .knowhere_hybrid import PersistedScoreCorpus + from shared.services.retrieval.scoring.knowhere_hybrid import PersistedScoreCorpus + + +@dataclass +class Chunk: + node_id: str + doc_id: str + text: str + line_ids: Tuple[int, ...] + section_id: Optional[str] = None + text_line_id_groups: Optional[Tuple[Tuple[int, ...], ...]] = None @dataclass @@ -107,10 +117,6 @@ class ProviderToolSpace: def __init__(self, provider: HierarchyProvider) -> None: self._provider = provider - def address_level(self, node_id: str): - fn = getattr(self._provider, "address_level", None) - return fn(node_id) if callable(fn) else None - def owner_document(self, node_id: str) -> Optional[str]: fn = getattr(self._provider, "owner_document", None) if not callable(fn): @@ -123,7 +129,7 @@ def document_ids(self) -> List[str]: fn = getattr(self._provider, "document_ids", None) if not callable(fn): return [] - return [str(x) for x in (fn() or ()) if str(x).strip()] + return [str(x) for x in cast(Sequence[Any], fn() or ()) if str(x).strip()] def sections_for_doc(self, doc_id: str) -> List[str]: return [str(s) for s in self._provider.roots(doc_id)] @@ -190,7 +196,7 @@ def _node_unit_span(self, section_id: str) -> Tuple[str, int, int]: unit_text = getattr(self._provider, "unit_text", None) if not callable(self_units) or not callable(unit_text): return "", 0, 0 - units = list(self_units(section_id) or ()) + units = list(cast(Sequence[Any], self_units(section_id) or ())) if not units: return "", 0, 0 first_order = int(getattr(units[0], "sort_order", 0) or 0) @@ -251,8 +257,6 @@ def _node_unit_span(self, section_id: str) -> Tuple[str, int, int]: def _make_chunk( self, node_id: str, doc_id: str, text: str, order: int, section_id: str ) -> Any: - from ._compat import Chunk # type: ignore - return Chunk( node_id=node_id, doc_id=doc_id, @@ -267,7 +271,7 @@ def materialize_self_only_chunks(self, section_id: str, doc_id: str) -> List[Any if not callable(self_units) or not callable(unit_text): return [] out: List[Any] = [] - for unit in self_units(section_id) or (): + for unit in cast(Sequence[Any], self_units(section_id) or ()): text = str(unit_text(unit) or "").strip() if not text: continue @@ -293,9 +297,9 @@ def _materialize_leaf_path_chunks(self, section_id: str, doc_id: str) -> List[An ] # One unit per descendant leaf, plus one per interstitial parent, so - # node ids line up with the keys nav_map_scores.build_score_units emits. + # node ids line up with the keys scoring.score_units.build_score_units emits. out: List[Any] = [] - for leaf_id in leaf_fn(section_id) or (): + for leaf_id in cast(Sequence[Any], leaf_fn(section_id) or ()): text, order, _count = self._node_unit_span(leaf_id) if text: out.append(self._make_chunk(leaf_id, doc_id, text, order, leaf_id)) @@ -320,102 +324,3 @@ def load_persisted_score_corpus( return None return cast(Optional["PersistedScoreCorpus"], fn(doc_ids, queries)) - -@dataclass -class InMemoryNode: - section_id: str - title: str - content: str = "" - children: List[str] = field(default_factory=list) - - -class InMemoryHierarchyProvider: - """Minimal reference ``HierarchyProvider``: no scoring, no ToolSpace. - - Built directly from a ``{doc_id: [InMemoryNode, ...]}`` map plus a - ``{doc_id: [root_section_id, ...]}`` map — the "hierarchy + summary is - enough" claim's simplest possible witness. - """ - - def __init__( - self, - *, - roots_by_doc: Dict[str, Sequence[str]], - nodes: Dict[str, InMemoryNode], - summaries: Optional[Dict[str, str]] = None, - ) -> None: - self._roots_by_doc = {k: list(v) for k, v in roots_by_doc.items()} - self._nodes = dict(nodes) - self._summaries = dict(summaries or {}) - self._parent: Dict[str, str] = {} - for node in self._nodes.values(): - for child_id in node.children: - self._parent[child_id] = node.section_id - self._owner: Dict[str, str] = {} - for doc_id, root_ids in self._roots_by_doc.items(): - stack = list(root_ids) - while stack: - sid = stack.pop() - if sid in self._owner: - continue - self._owner[sid] = doc_id - node = self._nodes.get(sid) - if node: - stack.extend(node.children) - - def owner_document(self, node_id: str) -> Optional[str]: - return self._owner.get(str(node_id or "").strip()) - - def roots(self, doc_id: str) -> Sequence[str]: - return list(self._roots_by_doc.get(doc_id, ())) - - def children(self, section_id: str) -> Sequence[str]: - node = self._nodes.get(section_id) - return list(node.children) if node else [] - - def node_meta(self, section_id: str) -> NodeMeta: - node = self._nodes.get(section_id) - if node is None: - return NodeMeta() - return NodeMeta( - title=node.title, - summary=self._summaries.get(section_id, ""), - has_children=bool(node.children), - ) - - def parent_id(self, section_id: str) -> Optional[str]: - return self._parent.get(section_id) - - def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: - ancestors: Set[str] = set() - cur = self._parent.get(section_id) - while cur: - ancestors.add(cur) - cur = self._parent.get(cur) - descendants: Set[str] = set() - stack = list(self.children(section_id)) - while stack: - cid = stack.pop() - if cid in descendants: - continue - descendants.add(cid) - stack.extend(self.children(cid)) - return ancestors, descendants - - def content(self, section_id: str) -> str: - node = self._nodes.get(section_id) - if node is None: - return "" - parts: List[str] = [] - - def walk(sid: str) -> None: - cur = self._nodes.get(sid) - if cur is None: - return - if cur.content: - parts.append(cur.content) - for cid in cur.children: - walk(cid) - - walk(section_id) - return "\n".join(parts) diff --git a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py rename to packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py index 6e5a95f8..ca73845b 100644 --- a/packages/shared-python/shared/services/retrieval/nav/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py @@ -247,7 +247,7 @@ def empty(cls) -> "_StreamingBm25Stats": def score( self, document_length: int, - frequencies: Dict[str, int], + frequencies: Mapping[str, int], query_tokens: List[str], ) -> float: if ( diff --git a/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py b/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py new file mode 100644 index 00000000..34c3b5c6 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoring/knowhere_provider.py @@ -0,0 +1,342 @@ +"""Publication-time hierarchy provider over section/chunk rows. + +Extracted from the map-nav package so document publication and classic +recall can build map-unit indexes without the episode kernel. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Set, + Tuple, +) + +from shared.services.retrieval.scoring.hierarchy import NodeMeta + + +_ASSET_TYPES = ("table", "image") +# Body chunk types that can own a Root-parked asset via connect_to. Both +# chunk-track ("text") and page-track ("page") body chunks can embed assets. +_BODY_CHUNK_TYPES = ("text", "page") +# Knowhere sentinel path for the virtual document container (not a collectable leaf). +ROOT_SECTION_PATH = "Root" + + +@dataclass(frozen=True) +class SectionRow: + """One ``document_sections`` row.""" + + section_id: str + parent_section_id: Optional[str] + section_path: str + section_title: str + section_level: int + summary: str + sort_order: int + + +@dataclass(frozen=True) +class UnitRow: + """One ``document_chunks`` row.""" + + chunk_id: str + section_id: Optional[str] + chunk_type: str + content: str + sort_order: int + source_chunk_path: str = "" + file_path: str = "" + metadata: Dict[str, Any] = field(default_factory=dict) + + +def asset_display_text(unit: UnitRow) -> str: + """Body text for an asset unit, whose ``content`` is only a file path. + + Mirrors knowhere's own assembly: an asset contributes its summary, not its + path. Without this an asset unit is unscorable and unreadable. + """ + meta = unit.metadata or {} + title = str(meta.get("asset_title") or "").strip() + summary = str(meta.get("summary") or "").strip() + ref = unit.file_path or unit.source_chunk_path or unit.content + label = "Table" if unit.chunk_type == "table" else "Image" + parts = [f"[{label}: {ref}]"] if ref else [f"[{label}]"] + if title: + parts.append(title) + if summary: + parts.append(summary) + return "\n".join(parts) + + +def normalize_section_path(path: str) -> str: + """Canonical path for gold/lookup: ``a / b`` (accepts ``a/b`` or ``a / b``).""" + raw = str(path or "").strip().strip("/") + if not raw or raw == ROOT_SECTION_PATH: + return "" + if " / " in raw: + parts = [p.strip() for p in raw.split(" / ") if p.strip()] + else: + parts = [p.strip() for p in raw.split("/") if p.strip()] + return " / ".join(parts) + + +def is_root_section_path(path: str) -> bool: + """True when the raw ``section_path`` is Knowhere's Root container.""" + return str(path or "").strip() == ROOT_SECTION_PATH + + +def _connect_to_targets(metadata: Dict[str, Any]) -> List[str]: + """``chunk_metadata.connect_to[].target`` ids (document order, first wins upstream).""" + raw = metadata.get("connect_to") if isinstance(metadata, dict) else None + if not isinstance(raw, list): + return [] + out: List[str] = [] + for conn in raw: + if not isinstance(conn, dict): + continue + target = str(conn.get("target") or "").strip() + if target: + out.append(target) + return out + + +class KnowhereProvider: + """``HierarchyProvider`` over knowhere section/chunk rows.""" + + def __init__( + self, + *, + doc_id: str, + sections: Sequence[SectionRow], + units: Sequence[UnitRow], + lazy_loader: Optional[Callable[[str], Sequence[UnitRow]]] = None, + known_chunk_ids: Optional[Sequence[str]] = None, + ) -> None: + self.doc_id = str(doc_id) + self._lazy_loader = lazy_loader + self._loaded_sections: Set[str] = set() + self._sections: Dict[str, SectionRow] = {s.section_id: s for s in sections} + self._children: Dict[str, List[str]] = {} + self._roots: List[str] = [] + self._path_to_id: Dict[str, str] = {} + for row in sorted(sections, key=lambda s: (s.sort_order, s.section_id)): + parent = row.parent_section_id + if parent and parent in self._sections: + self._children.setdefault(parent, []).append(row.section_id) + else: + self._roots.append(row.section_id) + key = normalize_section_path(row.section_path) + if key: + self._path_to_id[key] = row.section_id + + self._units_by_section: Dict[str, List[UnitRow]] = {} + self._chunk_ids: Set[str] = set() + if known_chunk_ids: + self._chunk_ids.update( + str(chunk_id).strip() + for chunk_id in known_chunk_ids + if str(chunk_id).strip() + ) + for unit in sorted(units, key=lambda u: (u.sort_order, u.chunk_id)): + sid = unit.section_id + if not sid or sid not in self._sections: + continue + self._units_by_section.setdefault(sid, []).append(unit) + if unit.chunk_id: + self._chunk_ids.add(unit.chunk_id) + self._remount_root_assets() + + def _ensure_section_loaded(self, section_id: str) -> None: + if self._lazy_loader is None or section_id in self._loaded_sections: + return + loaded = list(self._lazy_loader(section_id) or ()) + self._loaded_sections.add(section_id) + if not loaded: + return + current = self._units_by_section.setdefault(section_id, []) + known = {unit.chunk_id for unit in current} + for unit in loaded: + if unit.chunk_id and unit.chunk_id not in known: + current.append(unit) + known.add(unit.chunk_id) + current.sort(key=lambda unit: (unit.sort_order, unit.chunk_id)) + + def _remount_root_assets(self) -> None: + """Reattach Root-FK image|table units to host sections via ``connect_to``. + + Aligns with Knowhere ``resolve_root_asset_owners``: assets whose FK still + points at Root are owned by the text chunk that lists them in + ``metadata.connect_to``. Unresolved Root assets leave the evidence surface. + """ + root_sids = [ + sid + for sid, row in self._sections.items() + if is_root_section_path(row.section_path) + ] + if not root_sids: + return + + root_assets: Dict[str, UnitRow] = {} + for sid in root_sids: + for unit in self._units_by_section.get(sid, ()): + if unit.chunk_type in _ASSET_TYPES and unit.chunk_id: + root_assets[unit.chunk_id] = unit + if not root_assets: + return + + owner_by_asset: Dict[str, str] = {} + for sid, units in self._units_by_section.items(): + row = self._sections.get(sid) + if row is None or is_root_section_path(row.section_path): + continue + for unit in units: + if unit.chunk_type not in _BODY_CHUNK_TYPES: + continue + for target in _connect_to_targets(unit.metadata or {}): + if target in root_assets and target not in owner_by_asset: + owner_by_asset[target] = sid + + touched_owners: Set[str] = set() + for chunk_id, owner_sid in owner_by_asset.items(): + unit = root_assets[chunk_id] + remounted = UnitRow( + chunk_id=unit.chunk_id, + section_id=owner_sid, + chunk_type=unit.chunk_type, + content=unit.content, + sort_order=unit.sort_order, + source_chunk_path=unit.source_chunk_path, + file_path=unit.file_path, + metadata=dict(unit.metadata or {}), + ) + self._units_by_section.setdefault(owner_sid, []).append(remounted) + touched_owners.add(owner_sid) + + for sid in root_sids: + self._units_by_section[sid] = [ + u + for u in self._units_by_section.get(sid, ()) + if u.chunk_type not in _ASSET_TYPES + ] + for sid in touched_owners: + self._units_by_section[sid].sort(key=lambda u: (u.sort_order, u.chunk_id)) + + def owner_document(self, node_id: str) -> Optional[str]: + sid = str(node_id or "").strip() + if not sid: + return None + if sid == self.doc_id or sid in self._sections or sid in self._chunk_ids: + return self.doc_id + return None + + def roots(self, doc_id: str) -> Sequence[str]: + return list(self._roots) if str(doc_id) == self.doc_id else [] + + def children(self, section_id: str) -> Sequence[str]: + return list(self._children.get(section_id, ())) + + def node_meta(self, section_id: str) -> NodeMeta: + row = self._sections.get(section_id) + if row is None: + return NodeMeta() + return NodeMeta( + title=row.section_title, + summary=row.summary, + has_children=bool(self._children.get(section_id)), + ) + + def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: + ancestors: Set[str] = set() + cur = self._sections.get(section_id) + while cur is not None and cur.parent_section_id: + parent = cur.parent_section_id + if parent in ancestors: + break + ancestors.add(parent) + cur = self._sections.get(parent) + descendants: Set[str] = set() + stack = list(self.children(section_id)) + while stack: + cid = stack.pop() + if cid in descendants: + continue + descendants.add(cid) + stack.extend(self.children(cid)) + return ancestors, descendants + + def content(self, section_id: str) -> str: + units = self.subtree_units(section_id) + return "\n".join(self.unit_text(u) for u in units if self.unit_text(u)) + + def self_units(self, section_id: str) -> List[UnitRow]: + self._ensure_section_loaded(section_id) + return list(self._units_by_section.get(section_id, ())) + + def subtree_units(self, section_id: str) -> List[UnitRow]: + out = list(self.self_units(section_id)) + for cid in self.relations(section_id)[1]: + out.extend(self.self_units(cid)) + out.sort(key=lambda u: (u.sort_order, u.chunk_id)) + return out + + def leaf_ids(self, section_id: str) -> List[str]: + out: List[str] = [] + + def rec(sid: str) -> None: + kids = self.children(sid) + if not kids: + out.append(sid) + return + for kid in kids: + rec(kid) + + rec(section_id) + return out + + def path_titles(self, section_id: str) -> str: + chain: List[str] = [] + cur = self._sections.get(section_id) + while cur is not None: + if cur.section_title: + chain.append(cur.section_title) + parent = cur.parent_section_id + cur = self._sections.get(parent) if parent else None + return " / ".join(reversed(chain)) + + def parent_id(self, section_id: str) -> Optional[str]: + row = self._sections.get(section_id) + return row.parent_section_id if row else None + + def section_path(self, section_id: str) -> str: + row = self._sections.get(section_id) + return str(row.section_path or "") if row else "" + + def resolve_path(self, path: str) -> Optional[str]: + """Map a human/gold path to ``section_id`` (``sec_*``).""" + key = normalize_section_path(path) + if not key: + return None + return self._path_to_id.get(key) + + def unit_text(self, unit: UnitRow) -> str: + if unit.chunk_type in _ASSET_TYPES: + return asset_display_text(unit) + return str(unit.content or "").strip() + + def summaries(self) -> Dict[str, str]: + return { + sid: row.summary + for sid, row in self._sections.items() + if str(row.summary or "").strip() + } + + def all_section_ids(self) -> List[str]: + return list(self._sections) + diff --git a/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py b/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py new file mode 100644 index 00000000..91bfa047 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoring/node_filter_predicates.py @@ -0,0 +1,83 @@ +"""Predicate compile/match for section path/summary filters. + +Extracted from the map-nav tree walker. Live consumers evaluate these +predicates against persisted section rows, not an in-memory episode tree. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any, List, Literal, Sequence, Tuple + +MatchKind = Literal["substring", "regex"] +FilterField = Literal["path", "summary"] + +_MAX_REGEX_PATTERN_LEN = 256 + + +@dataclass(frozen=True) +class FieldPredicate: + field: FilterField + terms: Tuple[str, ...] + match: MatchKind = "substring" + + +def field_predicate( + field: str, + terms: Sequence[str], + match: str = "substring", +) -> FieldPredicate: + key = str(field or "").strip().lower() + if key not in {"path", "summary"}: + raise ValueError(f"unsupported filter field: {field!r}") + kind = str(match or "substring").strip().lower() + if kind not in {"substring", "regex"}: + raise ValueError(f"unsupported filter match: {match!r}") + cleaned = tuple(str(term) for term in terms if str(term)) + return FieldPredicate(field=key, terms=cleaned, match=kind) # type: ignore[arg-type] + + +def _compile_predicates( + predicates: Sequence[FieldPredicate], +) -> Tuple[List[Tuple[FieldPredicate, List[Any]]], List[str]]: + compiled: List[Tuple[FieldPredicate, List[Any]]] = [] + failed: List[str] = [] + for pred in predicates: + if pred.match != "regex": + compiled.append((pred, [])) + continue + patterns: List[Any] = [] + ok = True + for term in pred.terms: + if len(term) > _MAX_REGEX_PATTERN_LEN: + failed.append(f"{pred.field}:regex:too_long") + ok = False + break + try: + patterns.append(re.compile(term, flags=re.IGNORECASE)) + except re.error: + failed.append(f"{pred.field}:regex:invalid") + ok = False + break + if ok: + compiled.append((pred, patterns)) + return compiled, failed + + +def _node_matches( + values: dict[str, str], + compiled: Sequence[Tuple[FieldPredicate, List[Any]]], +) -> bool: + if not compiled: + return True + for pred, patterns in compiled: + text = values.get(pred.field, "") + if pred.match == "regex": + if not patterns or not any(p.search(text or "") for p in patterns): + return False + continue + haystack = (text or "").lower() + if not pred.terms or not any(term.lower() in haystack for term in pred.terms): + return False + return True diff --git a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py b/packages/shared-python/shared/services/retrieval/scoring/persisted_score_load.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py rename to packages/shared-python/shared/services/retrieval/scoring/persisted_score_load.py index 712f23b5..53924b01 100644 --- a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py +++ b/packages/shared-python/shared/services/retrieval/scoring/persisted_score_load.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any -from shared.services.retrieval.nav.knowhere_hybrid import PersistedBm25Stats +from shared.services.retrieval.scoring.knowhere_hybrid import PersistedBm25Stats def average_idf_from_unit_dfs( diff --git a/packages/shared-python/shared/services/retrieval/scoring/score_units.py b/packages/shared-python/shared/services/retrieval/scoring/score_units.py new file mode 100644 index 00000000..e9c05b0a --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoring/score_units.py @@ -0,0 +1,188 @@ +"""Build persisted map-unit scoring rows from a hierarchy tool space.""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, cast + +from shared.services.retrieval.scoring.knowhere_hybrid import ( + build_content_search_text, + build_path_search_text, + build_term_search_text, +) + +def _children_ids(ts: Any, section_id: str, doc_id: str) -> List[str]: + children_fn = getattr(ts, "_children_for_section_path", None) + if not callable(children_fn): + st = ts.get_structure(section_id) + rows = st.get("children") or [] + return [ + str(r.get("section_id") or "").strip() for r in rows if r.get("section_id") + ] + rows = cast(Sequence[Any], children_fn(section_id, doc_id)) + return [str(r.get("section_id") or "").strip() for r in rows if r.get("section_id")] + + +def _line_content(ts: Any, section_id: str, doc_id: str) -> str: + """Raw line text for a section node (no truncation).""" + idx = getattr(ts, "_idx", None) + b = getattr(idx, "_bundles", {}).get(doc_id) if idx is not None else None + if b is None: + path_fn = getattr(ts, "path_titles", None) + if callable(path_fn): + path = str(path_fn(section_id, doc_id) or "").strip() + return path.rsplit(" / ", 1)[-1] if path else "" + st = ts.get_structure(section_id) + return str(st.get("preview") or "").strip() + loc = getattr(idx, "_node_to_doc_line", {}).get(section_id) + if not loc: + return "" + _doc, line_idx = loc + if line_idx < 0 or line_idx >= len(b.lines): + return "" + return str(b.lines[line_idx].content or "").strip() + + +def _ancestor_path_titles(ts: Any, section_id: str, doc_id: str) -> str: + idx = getattr(ts, "_idx", None) + if idx is None: + # Provider-backed spaces expose the title chain directly; without this + # the path channel would score every unit as empty. + path_fn = getattr(ts, "path_titles", None) + return str(path_fn(section_id, doc_id) or "") if callable(path_fn) else "" + try: + ancestors = list(idx.ancestor_line_node_ids(section_id)) + except Exception: + ancestors = [] + titles: List[str] = [] + for aid in reversed(ancestors): + if not str(aid).startswith(f"{doc_id}:"): + continue + titles.append(_line_content(ts, aid, doc_id)) + titles.append(_line_content(ts, section_id, doc_id)) + return " / ".join(t for t in titles if t) + + +def _self_only_text(ts: Any, section_id: str, doc_id: str) -> Tuple[str, bool]: + """Return (self_text, has_interstitial_body). + + Interstitial means self_only span contains content beyond the heading line + itself (structural: more than one line/chunk in the self span). + """ + self_fn = getattr(ts, "materialize_self_only_chunks", None) + if not callable(self_fn): + return "", False + chunks = list(cast(Sequence[Any], self_fn(section_id, doc_id) or [])) + if not chunks: + return "", False + texts = [str(getattr(c, "text", "") or "").strip() for c in chunks] + texts = [t for t in texts if t] + if not texts: + return "", False + # Structural interstitial: self span covers more than the node heading line. + has_interstitial = len(chunks) > 1 + return "\n".join(texts), has_interstitial + + +def _section_body_text(ts: Any, section_id: str, doc_id: str) -> str: + """Heading + lines until first structural child (leaf body / parent self span).""" + text, _ = _self_only_text(ts, section_id, doc_id) + if text: + return text + return _line_content(ts, section_id, doc_id) + + +def _walk_tree( + ts: Any, + doc_id: str, + root_ids: Sequence[str], +) -> Tuple[Dict[str, List[str]], Set[str], Dict[str, str]]: + """Return children map, leaf ids, and title map for reachable nodes.""" + children_map: Dict[str, List[str]] = {} + titles: Dict[str, str] = {} + leaves: Set[str] = set() + seen: Set[str] = set() + + def walk(sid: str) -> None: + if not sid or sid in seen: + return + seen.add(sid) + titles[sid] = _line_content(ts, sid, doc_id) + kids = [c for c in _children_ids(ts, sid, doc_id) if c] + children_map[sid] = kids + if not kids: + leaves.add(sid) + return + for kid in kids: + walk(kid) + + for rid in root_ids: + walk(rid) + return children_map, leaves, titles + + +def build_score_units( + ts: Any, doc_id: str, root_ids: Optional[Sequence[str]] = None +) -> List[dict]: + """Build leaf (+ interstitial self_only) units for hybrid scoring.""" + if root_ids is None: + root_ids = list(ts.sections_for_doc(doc_id)) + children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) + units: List[dict] = [] + seen_unit_ids: Set[str] = set() + + for leaf_id in sorted(leaves): + content = _section_body_text(ts, leaf_id, doc_id) or ( + titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) + ) + path_text = _ancestor_path_titles(ts, leaf_id, doc_id) + unit_id = leaf_id + if unit_id in seen_unit_ids: + continue + seen_unit_ids.add(unit_id) + title = titles.get(leaf_id) or _line_content(ts, leaf_id, doc_id) + units.append( + { + "chunk_id": unit_id, + "section_id": leaf_id, + "kind": "leaf", + "content": content, + "path_text": path_text, + "path_search_text": build_path_search_text( + section_path=path_text, section_title=title or content + ), + "content_search_text": build_content_search_text(content), + "term_search_text": build_term_search_text( + content, path_text=path_text + ), + } + ) + + # Parents with interstitial self body. + for sid, kids in children_map.items(): + if not kids: + continue + self_text, has_interstitial = _self_only_text(ts, sid, doc_id) + if not has_interstitial or not self_text: + continue + unit_id = f"{sid}__self" + if unit_id in seen_unit_ids: + continue + seen_unit_ids.add(unit_id) + path_text = _ancestor_path_titles(ts, sid, doc_id) + units.append( + { + "chunk_id": unit_id, + "section_id": sid, + "kind": "self_only", + "content": self_text, + "path_text": path_text, + "path_search_text": build_path_search_text( + section_path=path_text, section_title=titles.get(sid) or "" + ), + "content_search_text": build_content_search_text(self_text), + "term_search_text": build_term_search_text( + self_text, path_text=path_text + ), + } + ) + return units diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 68e27402..b2a72978 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -1,7 +1,6 @@ """Classic-route discovery via the persisted map-unit BM25 scorer. -Replaces the retired chunk-level 3-channel SQL scan. Scoring uses the same -``score_persisted_corpus_many`` formula as map-nav (path + content only). +Scoring uses ``score_persisted_corpus_many`` over path + content only. ``chunk_types`` is optional: omitted means score every in-scope unit. When the request is image/table only, ``has_image`` / ``has_table`` (written at @@ -31,14 +30,14 @@ iter_connected_target_ids, normalize_chunk_type, ) -from shared.services.retrieval.nav.knowhere_hybrid import ( +from shared.services.retrieval.scoring.knowhere_hybrid import ( MAP_UNIT_INDEX_FORMAT_VERSION, PersistedScoreCorpus, PersistedScoreUnit, score_persisted_corpus_many, tokenize_query_for_ranker, ) -from shared.services.retrieval.nav.persisted_score_load import ( +from shared.services.retrieval.scoring.persisted_score_load import ( build_channel_bm25_stats, combine_average_idf, ) @@ -508,26 +507,11 @@ async def map_unit_discovery( ) except Exception as exc: logger.warning("retrieval index readiness publish failed: %s", exc) - logger.warning( - "retrieval map index incomplete user_id=%s namespace=%s " - "expected_revisions=%d indexed_revisions=%d fallback=legacy_fts", - user_id, - namespace, - len(expected_revisions), - len(index_parts), - ) - return await _legacy_chunk_discovery( - db, - 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 or [], - filter_mode=filter_mode, - revision_pins=revision_pins, + raise RuntimeError( + "retrieval map-unit index is incomplete or incompatible " + f"(user_id={user_id} namespace={namespace} " + f"expected_revisions={len(expected_revisions)} " + f"indexed_revisions={len(index_parts)})" ) if has_incomplete_index_statistics: logger.warning( @@ -701,103 +685,6 @@ async def map_unit_discovery( ) -async def _legacy_chunk_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, - signal_paths: list[str], - filter_mode: str, - revision_pins: Mapping[str, str] | None, -) -> DiscoveryResult: - """Bounded lexical fallback used while a serving index is incomplete.""" - clauses = [ - "d.user_id = :user_id", - "d.namespace = :namespace", - "d.status = 'active'", - ] - params: dict[str, Any] = { - "user_id": user_id, - "namespace": namespace, - "query": query, - "limit": max(1, int(top_k)), - } - if revision_pins is None: - clauses.append("d.current_job_result_id = dc.job_result_id") - else: - pairs = [ - (str(document_id), str(job_result_id)) - for document_id, job_result_id in revision_pins.items() - ] - if not pairs: - return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) - placeholders = [] - for index, (document_id, job_result_id) in enumerate(pairs): - document_key = f"_legacy_doc_{index}" - revision_key = f"_legacy_revision_{index}" - placeholders.append(f"(:{document_key}, :{revision_key})") - params[document_key] = document_id - params[revision_key] = job_result_id - clauses.append(f"(dc.document_id, dc.job_result_id) IN ({', '.join(placeholders)})") - if exclude_document_ids: - clauses.append("d.document_id <> ALL(:excluded_doc_ids)") - params["excluded_doc_ids"] = exclude_document_ids - if chunk_types: - type_keys = [] - for index, chunk_type in enumerate(sorted(chunk_types)): - key = f"_legacy_type_{index}" - type_keys.append(f":{key}") - params[key] = chunk_type - clauses.append(f"LOWER(dc.chunk_type) IN ({', '.join(type_keys)})") - if signal_paths: - signal_parts = [] - for index, signal in enumerate(signal_paths): - key = f"_legacy_signal_{index}" - signal_parts.append("LOWER(COALESCE(ds.section_path, '')) LIKE :" + key) - params[key] = f"%{signal.lower()}%" - combined = " OR ".join(signal_parts) - clauses.append(f"({combined})" if filter_mode == "keep" else f"NOT ({combined})") - for index, item in enumerate(exclude_sections): - document_id = str(item.get("document_id") or "").strip() - section_path = str(item.get("section_path") or "").strip() - if not document_id or not section_path: - continue - doc_key = f"_legacy_exclude_doc_{index}" - path_key = f"_legacy_exclude_path_{index}" - params[doc_key] = document_id - params[path_key] = section_path - clauses.append( - "NOT (dc.document_id = :" + doc_key + " AND (" - "COALESCE(ds.section_path, '') = :" + path_key + " OR " - "POSITION(:" + path_key + " || ' / ' IN COALESCE(ds.section_path, '')) = 1))" - ) - where_sql = " AND ".join(clauses) - statement = text( - "SELECT dc.chunk_id, dc.document_id, dc.section_id, dc.chunk_type, " - "dc.content, dc.source_chunk_path, dc.file_path, dc.chunk_metadata, " - "dc.job_result_id, dc.sort_order, ds.section_path, d.source_file_name, " - "jr.job_id, GREATEST(ts_rank_cd(dc.path_search_tsv, plainto_tsquery('simple', :query)), " - "2 * ts_rank_cd(dc.content_search_tsv, plainto_tsquery('simple', :query))) AS score " - "FROM document_chunks dc JOIN documents d ON d.document_id = dc.document_id " - "LEFT JOIN document_sections ds ON ds.section_id = dc.section_id " - "LEFT JOIN job_results jr ON jr.id = dc.job_result_id " - f"WHERE {where_sql} AND (dc.path_search_tsv @@ plainto_tsquery('simple', :query) " - "OR dc.content_search_tsv @@ plainto_tsquery('simple', :query) " - "OR LOWER(COALESCE(dc.term_search_text, '')) LIKE LOWER(:term_query)) " - "ORDER BY score DESC, dc.sort_order, dc.chunk_id LIMIT :limit" - ) - params["term_query"] = f"%{query}%" - rows = [dict(row._mapping) for row in (await db.execute(statement, params)).all()] - if rows: - normalize_row_scores(rows, source_field="score", target_field="discovery_score", default=0.5) - return DiscoveryResult(status="discovery_done", payload={"fused_rows": rows}) - - def _as_metadata_dict(value: object) -> dict[str, Any]: if isinstance(value, dict): return value diff --git a/packages/shared-python/shared/services/retrieval/serving_manifest.py b/packages/shared-python/shared/services/retrieval/serving_manifest.py index 77a33e1c..11b727ee 100644 --- a/packages/shared-python/shared/services/retrieval/serving_manifest.py +++ b/packages/shared-python/shared/services/retrieval/serving_manifest.py @@ -23,6 +23,10 @@ SERVING_MANIFEST_FORMAT_VERSION = 1 NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION = 2 +# Body chunk types that can own a Root-parked asset via connect_to. Both +# chunk-track ("text") and page-track ("page") body chunks can embed assets. +_BODY_CHUNK_TYPES = {"text", "page"} + def build_revision_serving_payload( db: Session, @@ -66,7 +70,9 @@ def build_revision_serving_payload( } remounted_assets: dict[str, list[str]] = {} for chunk in chunks: - if chunk.chunk_type != "text" or not isinstance(chunk.chunk_metadata, dict): + if chunk.chunk_type not in _BODY_CHUNK_TYPES or not isinstance( + chunk.chunk_metadata, dict + ): continue connections = chunk.chunk_metadata.get("connect_to") if not isinstance(connections, list): diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py index f2e24b56..4dd3130b 100644 --- a/packages/shared-python/shared/services/retrieval/settings.py +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -6,6 +6,10 @@ RRF_K = 60 DEFAULT_TOP_K = 10 +# Final evidence / tool-observation text budget (characters). Used by +# agent tool-loop harness caps (``ToolBudget.max_chars``). +EVIDENCE_TEXT_CHAR_BUDGET = 12_000 + VALID_CHUNK_TYPES: set[str] = {"text", "image", "table", "page"} ASSET_CHUNK_TYPES: set[str] = {"image", "table"} diff --git a/packages/shared-python/shared/services/retrieval/trace/__init__.py b/packages/shared-python/shared/services/retrieval/trace/__init__.py index 54152432..657a5ebe 100644 --- a/packages/shared-python/shared/services/retrieval/trace/__init__.py +++ b/packages/shared-python/shared/services/retrieval/trace/__init__.py @@ -1,21 +1,9 @@ """Retrieval decision-trace package (replaces agentic/core trace types).""" -from shared.services.retrieval.trace.mapnav import ( - build_decision_trace, - episode_selected_doc_ids, - episode_selected_paths, - episode_token_count, - episode_workflow_plan, -) from shared.services.retrieval.trace.recorder import TraceRecorder from shared.services.retrieval.trace.types import DecisionTraceStep __all__ = [ "DecisionTraceStep", "TraceRecorder", - "build_decision_trace", - "episode_workflow_plan", - "episode_selected_paths", - "episode_selected_doc_ids", - "episode_token_count", ] diff --git a/packages/shared-python/shared/services/retrieval/trace/recorder.py b/packages/shared-python/shared/services/retrieval/trace/recorder.py index 79d5cea7..8b5156a3 100644 --- a/packages/shared-python/shared/services/retrieval/trace/recorder.py +++ b/packages/shared-python/shared/services/retrieval/trace/recorder.py @@ -15,7 +15,6 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.nav_config import MAPNAV_MODEL from shared.services.retrieval.settings import DEFAULT_TOP_K from shared.services.retrieval.trace.types import DecisionTraceStep @@ -47,7 +46,7 @@ def __init__( policy_name: str = "llm_policy_v1", config: Any = None, ) -> None: - del config # legacy AgentRunConfig; ignored on map-nav path + del config # unused AgentRunConfig leftover; ignored self._db = db self._run_id = f"aret_{uuid4().hex[:12]}" self._user_id = user_id @@ -236,7 +235,7 @@ async def complete( "router": router_used, "step_count": len(self._steps), "final_doc_ids": doc_ids_in_result, - "model_name": model_name or MAPNAV_MODEL, + "model_name": model_name or "unknown", } if selected_paths is not None: provenance["selected_paths"] = selected_paths diff --git a/packages/shared-python/shared/tests/test_agent_explore_harness.py b/packages/shared-python/shared/tests/test_agent_explore_harness.py new file mode 100644 index 00000000..7b9740ce --- /dev/null +++ b/packages/shared-python/shared/tests/test_agent_explore_harness.py @@ -0,0 +1,195 @@ +"""Unit tests for the Phase 3.5 ``agent_explore`` harness layer. + +Covers only the pure functions and the ``AGENT_EXPLORE_HARNESS`` switch +resolution — none of these need a real DB or LLM. Does not cover +``dispatch.dispatch_tool_call`` (needs a DB session) or a full +``Harness.run_episode`` loop (needs an LLM/Cursor SDK backend) — those stay +integration-level, exercised via the debug scripts and +``eval-cursor-harness``, not here. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import pytest + +from shared.services.retrieval.agent_explore.harness.base import Harness +from shared.services.retrieval.agent_explore.harness.resolve import ( + _HARNESS_ENV, + resolve_harness, + resolve_harness_name, +) +from shared.services.retrieval.agent_explore.shared import ( + EVIDENCE_TOOL_NAMES, + build_wire_tool_name_map, + dedup_refs, + normalize_finish_refs, + tool_message_content, + wire_safe_tool_name, +) +from shared.services.retrieval.agent_tools import ToolResult + + +# -------------------------------------------------------------------------- +# shared.py: wire-safe tool name mapping +# -------------------------------------------------------------------------- + + +def test_wire_safe_tool_name_replaces_dots() -> None: + assert wire_safe_tool_name("corpus.read") == "corpus_read" + assert wire_safe_tool_name("corpus.node_filter") == "corpus_node_filter" + # Already wire-safe names are untouched. + assert wire_safe_tool_name("finish") == "finish" + + +def test_build_wire_tool_name_map_round_trips_to_canonical() -> None: + mapping = build_wire_tool_name_map(["corpus.read", "corpus.recall", "finish"]) + assert mapping == { + "corpus_read": "corpus.read", + "corpus_recall": "corpus.recall", + "finish": "finish", + } + + +# -------------------------------------------------------------------------- +# shared.py: tool_message_content (max_chars capping) +# -------------------------------------------------------------------------- + + +def test_tool_message_content_passes_through_short_text() -> None: + result = ToolResult(text="short body") + assert tool_message_content(result, max_chars=100) == "short body" + + +def test_tool_message_content_caps_long_text_with_note() -> None: + result = ToolResult(text="x" * 200) + content = tool_message_content(result, max_chars=100) + assert content.startswith("x" * 100) + assert "truncated, 100 more chars" in content + assert len(content) > 100 # capped body + truncation note, not silently dropped + + +def test_tool_message_content_surfaces_error_instead_of_text() -> None: + result = ToolResult(text="ignored", error="bad args: missing document_id") + assert tool_message_content(result, max_chars=100) == "error: bad args: missing document_id" + + +def test_tool_message_content_empty_text_placeholder() -> None: + result = ToolResult(text="") + assert tool_message_content(result, max_chars=100) == "(empty result)" + + +# -------------------------------------------------------------------------- +# shared.py: normalize_finish_refs / dedup_refs +# -------------------------------------------------------------------------- + + +def test_normalize_finish_refs_drops_non_dict_and_empty_document_id() -> None: + raw = [ + {"document_id": "doc_a", "chunk_id": "c1"}, + {"document_id": "", "chunk_id": "c2"}, + {"chunk_id": "c3"}, + "not a dict", + None, + {"document_id": "doc_b"}, + ] + normalized = normalize_finish_refs(raw) + assert normalized == [ + {"document_id": "doc_a", "chunk_id": "c1"}, + {"document_id": "doc_b"}, + ] + + +def test_normalize_finish_refs_non_list_input_is_empty() -> None: + assert normalize_finish_refs(None) == [] + assert normalize_finish_refs("refs") == [] + assert normalize_finish_refs({"document_id": "doc_a"}) == [] + + +def test_dedup_refs_keeps_first_seen_and_drops_missing_ids() -> None: + refs = [ + {"document_id": "doc_a", "chunk_id": "c1", "section_path": "first"}, + {"document_id": "doc_a", "chunk_id": "c1", "section_path": "duplicate"}, + {"document_id": "doc_a", "chunk_id": "c2"}, + {"document_id": "doc_a"}, # missing chunk_id -> dropped + {"chunk_id": "c3"}, # missing document_id -> dropped + {"document_id": "doc_b", "chunk_id": "c1"}, + ] + deduped = dedup_refs(refs) + assert deduped == [ + {"document_id": "doc_a", "chunk_id": "c1", "section_path": "first"}, + {"document_id": "doc_a", "chunk_id": "c2"}, + {"document_id": "doc_b", "chunk_id": "c1"}, + ] + + +def test_evidence_tool_names_is_read_and_assets_only() -> None: + assert EVIDENCE_TOOL_NAMES == frozenset({"corpus.read", "corpus.assets"}) + + +# -------------------------------------------------------------------------- +# harness/resolve.py: AGENT_EXPLORE_HARNESS switch +# -------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _clean_harness_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(_HARNESS_ENV, raising=False) + + +def test_resolve_harness_name_defaults_to_cursor_sdk() -> None: + assert resolve_harness_name() == "cursor_sdk" + + +def test_resolve_harness_name_reads_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(_HARNESS_ENV, "cursor_sdk") + assert resolve_harness_name() == "cursor_sdk" + + +def test_resolve_harness_name_unknown_value_falls_back_to_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(_HARNESS_ENV, "not_a_real_harness") + assert resolve_harness_name() == "cursor_sdk" + + +def test_resolve_harness_name_is_case_insensitive(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(_HARNESS_ENV, "CURSOR_SDK") + assert resolve_harness_name() == "cursor_sdk" + + +def test_resolve_harness_default_builds_cursor_harness() -> None: + from shared.services.retrieval.agent_explore.harness.cursor_harness import CursorHarness + + harness = resolve_harness() + assert isinstance(harness, CursorHarness) + assert isinstance(harness, Harness) + + +def test_resolve_harness_cursor_sdk_builds_cursor_harness_without_sdk_installed() -> None: + """Selecting cursor_sdk must not require the optional cursor-sdk package + to be importable — only actually running an episode does (see + cursor_harness.py's guarded _require_cursor_sdk, exercised at + run_episode() call time, not at harness construction time). + """ + from shared.services.retrieval.agent_explore.harness.cursor_harness import CursorHarness + + harness = resolve_harness("cursor_sdk") + assert isinstance(harness, CursorHarness) + assert isinstance(harness, Harness) + + +def test_resolve_harness_explicit_name_overrides_env(monkeypatch: pytest.MonkeyPatch) -> None: + from shared.services.retrieval.agent_explore.harness.openai_harness import OpenAIHarness + + monkeypatch.setenv(_HARNESS_ENV, "cursor_sdk") + harness = resolve_harness("openai") + assert isinstance(harness, OpenAIHarness) diff --git a/packages/shared-python/shared/tests/test_asset_inline.py b/packages/shared-python/shared/tests/test_asset_inline.py index a59652f0..d8345c24 100644 --- a/packages/shared-python/shared/tests/test_asset_inline.py +++ b/packages/shared-python/shared/tests/test_asset_inline.py @@ -10,8 +10,8 @@ from shared.services.retrieval.hydration.result_assembly import ( assemble_retrieval_results, ) -from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace -from shared.services.retrieval.nav.nav_knowhere import ( +from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace +from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, SectionRow, UnitRow, diff --git a/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py b/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py index 2d4b96e1..8b559ebb 100644 --- a/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py +++ b/packages/shared-python/shared/tests/test_knowhere_hybrid_tokenize.py @@ -2,7 +2,7 @@ from __future__ import annotations -from shared.services.retrieval.nav.knowhere_hybrid import ( +from shared.services.retrieval.scoring.knowhere_hybrid import ( MAP_UNIT_INDEX_FORMAT_VERSION, build_content_search_text, build_path_search_text, diff --git a/packages/shared-python/shared/tests/test_section_path_lookup.py b/packages/shared-python/shared/tests/test_section_path_lookup.py new file mode 100644 index 00000000..037d7e1f --- /dev/null +++ b/packages/shared-python/shared/tests/test_section_path_lookup.py @@ -0,0 +1,39 @@ +"""Pure tests for agent_tools section_path suffix resolution.""" + +from __future__ import annotations + +from shared.services.retrieval.agent_tools.section_path_lookup import ( + format_ambiguous_section_path_error, + paths_matching_section_ref, +) + + +def test_paths_matching_section_ref_exact() -> None: + paths = ["附件目录 / 1.1 概述", "综合说明 / 1.1 概述"] + assert paths_matching_section_ref("附件目录 / 1.1 概述", paths) == ["附件目录 / 1.1 概述"] + + +def test_paths_matching_section_ref_suffix() -> None: + paths = [ + "附件目录 / 3 工程地质 / 3.2 覆盖层", + "综合说明 / 3 工程地质 / 3.2 覆盖层", + ] + assert paths_matching_section_ref("3 工程地质 / 3.2 覆盖层", paths) == sorted(paths) + + +def test_paths_matching_section_ref_no_substring_false_positive() -> None: + paths = ["附件目录 / 3.2 覆盖层处理", "附件目录 / 13.2 覆盖层"] + assert paths_matching_section_ref("3.2 覆盖层", paths) == [] + + +def test_paths_matching_section_ref_top_level() -> None: + paths = ["附件目录", "附件目录 / 1.1 概述"] + assert paths_matching_section_ref("附件目录", paths) == ["附件目录"] + + +def test_format_ambiguous_section_path_error_lists_candidates() -> None: + matches = ["A / X", "B / X"] + msg = format_ambiguous_section_path_error("X", matches) + assert "ambiguous section_path 'X'" in msg + assert "A / X" in msg + assert "B / X" in msg diff --git a/pyproject.toml b/pyproject.toml index 19876076..b764a83e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ exclude = [ "packages/shared-python/build", "packages/shared-python/dist", "packages/shared-python/shared/tests", - "packages/shared-python/shared/services/retrieval/nav", + "deprecated", ] executionEnvironments = [ { root = "apps/api", extraPaths = ["packages/shared-python"] }, @@ -63,6 +63,7 @@ src = [ "apps/worker", "packages/shared-python", ] +exclude = ["deprecated"] [tool.ruff.lint] select = ["E", "F"] diff --git a/pytest.ini b/pytest.ini index e83f7c02..cd7e4782 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,4 +1,6 @@ [pytest] +norecursedirs = deprecated .* +addopts = --ignore=deprecated asyncio_mode = auto asyncio_default_fixture_loop_scope = function filterwarnings = diff --git a/uv.lock b/uv.lock index 77a0eee5..2e598298 100644 --- a/uv.lock +++ b/uv.lock @@ -805,6 +805,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/20/2a/1b016902351a523aa2bd446b50a5bc1175d7a7d1cf90fe2ef904f9b84ebc/cryptography-46.0.7-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:258514877e15963bd43b558917bc9f54cf7cf866c38aa576ebf47a77ddbc43a4", size = 3412829, upload-time = "2026-04-08T01:57:48.874Z" }, ] +[[package]] +name = "cursor-sdk" +version = "1.0.31" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/82/ca/60a7ea6a8b08a430a499797261487cf9114fbbfb652e10379a1f98dde463/cursor_sdk-1.0.31.tar.gz", hash = "sha256:fcdd279852d0b3eea4e4c4562dcd1c7d18360f507d2830bd2de79b6d855276a5", size = 1177, upload-time = "2026-09-03T21:04:04.352Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/18/1ea7bd9823dde860a44c68a0e920d07a415e5f73c4c3693cc8f22a081220/cursor_sdk-1.0.31-py3-none-macosx_11_0_arm64.whl", hash = "sha256:0b7b1e8fb677fe400c9d63518ff08465a7527cd23f20a3d857bf9454ab99c411", size = 49516626, upload-time = "2026-09-03T21:03:49.927Z" }, + { url = "https://files.pythonhosted.org/packages/ca/21/df0f17f3bd3d956897756e5bf4dbed355ad9261d0b7efa3d1af703912fee/cursor_sdk-1.0.31-py3-none-macosx_11_0_x86_64.whl", hash = "sha256:22cb11230491cd989eb7015cb3778f09ba1e6a974661ddace5bee1e68f85ee8b", size = 51027071, upload-time = "2026-09-03T21:03:52.837Z" }, + { url = "https://files.pythonhosted.org/packages/28/96/de1b8b0275c792d9a472b8e61fe12e0d9bcd3f49551658361f43d8d381ac/cursor_sdk-1.0.31-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ae97e49570887953922eb6b3b84cb02c2d13e795fe4b3968f25e23ab569cf21f", size = 58623238, upload-time = "2026-09-03T21:03:55.741Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f1/25188d01c7bd0b42edee594225353176f90747047c1508bc315c0d2a3706/cursor_sdk-1.0.31-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9b87cd80fe428b1a2ee88b4a08aac6297179a33f809d650f0383901fb60a766", size = 59314084, upload-time = "2026-09-03T21:03:58.675Z" }, + { url = "https://files.pythonhosted.org/packages/de/60/e36b7d68ae702781fac63fef564fc9d643495b114fb959e442d449e537cf/cursor_sdk-1.0.31-py3-none-win_amd64.whl", hash = "sha256:12d87c639ac1bdb50958028e81f295f856e8e9a4d784977ea0e07d196716410e", size = 45610437, upload-time = "2026-09-03T21:04:01.878Z" }, +] + [[package]] name = "deprecated" version = "1.3.1" @@ -1397,6 +1413,7 @@ dependencies = [ { name = "aiohttp" }, { name = "alembic" }, { name = "celery" }, + { name = "cursor-sdk" }, { name = "fastapi" }, { name = "httpx" }, { name = "knowhere-shared" }, @@ -1431,6 +1448,7 @@ requires-dist = [ { name = "aiohttp", specifier = "==3.13.4" }, { name = "alembic", specifier = "==1.13.1" }, { name = "celery", specifier = "==5.5.3" }, + { name = "cursor-sdk", specifier = ">=1.0.31" }, { name = "fastapi", specifier = "==0.135.1" }, { name = "httpx", specifier = "==0.28.1" }, { name = "knowhere-shared", editable = "packages/shared-python" }, @@ -1588,6 +1606,11 @@ dependencies = [ { name = "tqdm" }, ] +[package.optional-dependencies] +cursor-harness = [ + { name = "cursor-sdk" }, +] + [package.dev-dependencies] dev = [ { name = "fakeredis", extra = ["lua"] }, @@ -1601,6 +1624,7 @@ dev = [ requires-dist = [ { name = "beautifulsoup4", specifier = "==4.13.4" }, { name = "cryptography", specifier = "==46.0.7" }, + { name = "cursor-sdk", marker = "extra == 'cursor-harness'", specifier = ">=1.0.31" }, { name = "gevent", specifier = ">=24.11.1" }, { name = "httpcore", specifier = ">=1.0.6" }, { name = "jieba", specifier = "==0.42.1" }, @@ -1624,6 +1648,7 @@ requires-dist = [ { name = "tabula-py", specifier = ">=2.10.0" }, { name = "tqdm", specifier = "==4.67.1" }, ] +provides-extras = ["cursor-harness"] [package.metadata.requires-dev] dev = [