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..1b961d4f 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -20260506-102713 +Knowhere 2.0

Prepare unstructured data for AI Agents

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..3a52b2fb --- /dev/null +++ b/apps/api/app/mcp/dynamic_tools.py @@ -0,0 +1,80 @@ +"""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 +from shared.services.retrieval.agent_tools import tools as _agent_tools_registered # noqa: F401 + +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/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..752ed091 --- /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/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/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..45c236f7 --- /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. + """ + ... 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..0d84b566 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py @@ -0,0 +1,303 @@ +"""``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, +) +from shared.services.retrieval.agent_tools import tools as _agent_tools_registered # noqa: F401 + +# 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..2d48f505 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py @@ -0,0 +1,390 @@ +"""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 +from shared.services.retrieval.agent_tools import tools as _agent_tools_registered # noqa: F401 + +# 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] = [] + stop_reason = "finished" + 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..9f5b4284 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py @@ -0,0 +1,39 @@ +"""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. +""" + +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 + +__all__ = [ + "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 = [