diff --git a/AGENTS.md b/AGENTS.md index 2bf644d2d..43f64487c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -576,9 +576,9 @@ flowchart LR **Channel weights** (default): path=1.0, content=2.0. **RRF formula**: `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 and are read by -`corpus.recall`'s term channel, not by classic discovery. +There is no scored term channel in this primary path. Chunk +`term_search_text` is persisted at publish time and is read by +`corpus.grep`, not by classic discovery. #### Agent-explore Mode (default) diff --git a/CONTEXT.md b/CONTEXT.md index 8e75965fd..5b465357f 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -343,7 +343,7 @@ The query workflow that returns cited evidence from published documents. ### Retrieval Query The typed retrieval request that owns cache-shaping fields and route policy: -scope, filters, data type, channels, ranking options, and agentic toggle. +scope, filters, data type, ranking options, and agentic toggle. ### Retrieval Run diff --git a/apps/api/.env.example b/apps/api/.env.example index f5a003ee2..bba577426 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -136,6 +136,12 @@ MOESIF_APPLICATION_ID= FREE_PLAN_INITIAL_CREDITS=5 FRONTEND_URL=http://localhost:3000 +# Reuse canonical demo bundles; disable to restore per-job uploads while keeping +# existing canonical artifact URLs readable. Does not change publication strategy. +DEMO_CANONICAL_BUNDLE_ENABLED=false +# Optional per-process tokenization/frequency reuse for versioned demo sources. +DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=false + # Required for specific features: parsing providers MINERU_API_KEYS= ILOVEAPI_PUBLIC_KEY= diff --git a/apps/api/alembic/versions/2b3c4d5e6f7a_drop_map_unit_term_search_text.py b/apps/api/alembic/versions/2b3c4d5e6f7a_drop_map_unit_term_search_text.py new file mode 100644 index 000000000..7936d0fb3 --- /dev/null +++ b/apps/api/alembic/versions/2b3c4d5e6f7a_drop_map_unit_term_search_text.py @@ -0,0 +1,55 @@ +"""Drop unused map-unit term-search leftovers. + +Classic scoring is path + content only. ``term_search_text_lower`` on +``document_map_units`` and its trigram index had no reader. Chunk +``term_search_text`` and ``idx_document_chunks_term_trgm`` stay: grep uses them. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "2b3c4d5e6f7a" +down_revision: str | None = "2f3g4h5i6j7" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +__all__ = [ + "revision", + "down_revision", + "branch_labels", + "depends_on", + "upgrade", + "downgrade", +] + +_MAP_UNIT_INDEX = "idx_document_map_units_term_trgm" + + +def upgrade() -> None: + op.execute(f"DROP INDEX IF EXISTS {_MAP_UNIT_INDEX}") + inspector = sa.inspect(op.get_bind()) + columns = {col["name"] for col in inspector.get_columns("document_map_units")} + if "term_search_text_lower" in columns: + op.drop_column("document_map_units", "term_search_text_lower") + + +def downgrade() -> None: + inspector = sa.inspect(op.get_bind()) + columns = {col["name"] for col in inspector.get_columns("document_map_units")} + if "term_search_text_lower" not in columns: + op.add_column( + "document_map_units", + sa.Column("term_search_text_lower", sa.Text(), nullable=False, server_default=""), + ) + op.alter_column("document_map_units", "term_search_text_lower", server_default=None) + op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") + op.execute( + f"CREATE INDEX IF NOT EXISTS {_MAP_UNIT_INDEX} " + "ON document_map_units USING gin " + "(term_search_text_lower gin_trgm_ops)" + ) diff --git a/apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py b/apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py new file mode 100644 index 000000000..5d04d1313 --- /dev/null +++ b/apps/api/alembic/versions/3c4d5e6f7a8b_gate_legacy_token_lookup_retirement.py @@ -0,0 +1,56 @@ +"""Opt in to retiring the legacy token lookup after replacement validation.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import context, op + +from shared.services.retrieval.maintenance.legacy_token_lookup import ( + LegacyTokenLookupMaintenance, +) + +revision: str = "3c4d5e6f7a8b" +down_revision: str | None = "2b3c4d5e6f7a" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def _run_maintenance(*, is_restore: bool) -> None: + if context.is_offline_mode(): + raise RuntimeError("Legacy token index maintenance requires live catalog checks") + is_external_transaction: bool = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + + def execute() -> None: + maintenance: LegacyTokenLookupMaintenance = LegacyTokenLookupMaintenance( + op.get_bind(), is_concurrent=not is_external_transaction + ) + if is_restore: + maintenance.restore_index() + else: + maintenance.apply_retirement() + + if is_external_transaction: + execute() + return + with op.get_context().autocommit_block(): + execute() + + +def upgrade() -> None: + """Default deployment records the revision and retains the legacy index.""" + gate: str | None = context.get_x_argument(as_dictionary=True).get( + "retire_legacy_token_lookup" + ) + if gate is None or gate == "false": + return + if gate != "true": + raise ValueError("retire_legacy_token_lookup must be exactly true or false") + _run_maintenance(is_restore=False) + + +def downgrade() -> None: + """Restore the exact legacy definition, including after operator retirement.""" + _run_maintenance(is_restore=True) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 1bd31d23b..b56185803 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -7,7 +7,7 @@ from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, field_validator from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db @@ -64,20 +64,6 @@ class RetrievalQueryRequest(BaseModel): filter_mode: Literal["delete", "keep"] = Field( "delete", description="Signal path filter mode" ) - channels: list[str] = Field( - default_factory=list, - description=( - "Deprecated and unsupported by the persisted map-unit route. " - "Leave empty; explicit channel selection is rejected." - ), - ) - channel_weights: dict[str, float] = Field( - default_factory=dict, - description=( - "Deprecated and unsupported by the persisted map-unit route. " - "Leave empty; explicit overrides are rejected." - ), - ) rerank: bool = Field(False, description="Enable LLM reranking after RRF fusion") threshold: float = Field(0.0, ge=0.0, description="Minimum RRF score threshold") internal_recall_k: int | None = Field( @@ -111,15 +97,6 @@ class RetrievalQueryRequest(BaseModel): ), ) - @field_validator("channels") - @classmethod - def validate_channels(cls, v: list[str]) -> list[str]: - valid = {"path", "content", "term"} - for ch in v: - if ch not in valid: - raise ValueError(f"Invalid channel: {ch}. Must be one of {valid}") - return v - @field_validator("chunk_types") @classmethod def validate_chunk_types(cls, v: list[str] | None) -> list[str] | None: @@ -137,20 +114,6 @@ def validate_chunk_types(cls, v: list[str] | None) -> list[str] | None: def normalize_namespace(cls, namespace: str | None) -> str: return normalize_retrieval_namespace(namespace) - @model_validator(mode="after") - def reject_unsupported_channel_controls(self) -> "RetrievalQueryRequest": - if self.channels: - raise ValueError( - "channels is deprecated and unsupported; omit it and use the " - "persisted path/content map-unit scorer" - ) - if self.channel_weights: - raise ValueError( - "channel_weights is deprecated and unsupported; omit it and use " - "the persisted path/content map-unit scorer" - ) - return self - class RetrievalQueryResponse(BaseModel): namespace: str @@ -209,11 +172,25 @@ async def execute_retrieval_query( else: resolved_chunk_types = None + query = str(payload.query or "").strip() + if not query: + return { + "namespace": normalize_retrieval_namespace(payload.namespace), + "query": query, + "router_used": "empty_query_filtered", + "failure_reason": "empty query — retrieval was not run", + "evidence": [], + "evidence_text": "", + "answer_text": "", + "referenced_chunks": [], + "results": [], + } + return await run_retrieval_query( db=db, user_id=current_user.user_id, namespace=normalize_retrieval_namespace(payload.namespace), - query=payload.query, + query=query, top_k=payload.top_k, include_document_ids=payload.include_document_ids, exclude_document_ids=payload.exclude_document_ids, @@ -221,8 +198,6 @@ async def execute_retrieval_query( chunk_types=resolved_chunk_types, signal_paths=payload.signal_paths or None, filter_mode=payload.filter_mode, - channels=payload.channels or None, - channel_weights=payload.channel_weights or None, rerank=payload.rerank, threshold=payload.threshold, internal_recall_k=payload.internal_recall_k, diff --git a/apps/api/app/mcp/dynamic_tools.py b/apps/api/app/mcp/dynamic_tools.py index 9586dfb43..c0b29607d 100644 --- a/apps/api/app/mcp/dynamic_tools.py +++ b/apps/api/app/mcp/dynamic_tools.py @@ -20,7 +20,7 @@ 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 pydantic import ConfigDict, create_model from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.agent_tools import REGISTRY, ToolContext, ToolSpec @@ -28,22 +28,65 @@ DbFactory = Callable[[], AsyncContextManager[AsyncSession]] +class _AllowExtraArgModel(ArgModelBase): + """``ArgModelBase`` that forwards only actually-supplied keys, plus any + unrecognized ones. + + ``FuncMetadata.call_fn_with_arg_validation`` builds ``_dispatch_tool``'s + kwargs by calling ``model_dump_one_level()``. The base implementation + (verified by reading ``mcp.server.fastmcp.utilities.func_metadata`` + directly) iterates every *declared* ``model_fields`` name unconditionally + — since every property is declared here as ``(Any, None)`` (see + ``_lenient_arg_model`` below), an MCP client that simply omits an + optional argument still gets that argument dumped as an explicit + ``None`` into ``REGISTRY.dispatch``'s args dict. ``validate_tool_args`` + then fails a perfectly legal call, e.g. ``corpus.grep(pattern="x")`` + reaching dispatch as ``{"pattern": "x", "chunk_types": None, ...}``, + which trips ``chunk_types``'s ``"type": "array"`` schema check. This was + a real, reproduced MCP-entry contract break, not a hypothetical one. + + Fix: dump only ``self.model_fields_set`` (pydantic's own record of which + declared fields were actually present in the validated input — verified + live that omitted fields are absent from this set even when every field + has a ``None`` default) plus ``model_extra`` for genuinely unknown keys. + Declaring ``extra="allow"`` alone is not enough for the *unknown-key* + half: pydantic stores those on the instance (``model_extra``) but the + base implementation never reads them; merging them in here is what gets + an unknown key from an MCP client into ``dispatch`` -> ``validate_tool_args``'s + unified unknown-argument error, instead of the previous fully-permissive + model silently dropping it before dispatch ever saw it. Together this + makes an MCP call's dumped kwargs exactly the set of keys the client + actually sent — omitted optional keys absent, unknown keys present (and + then rejected by schema validation), never a fabricated ``None``. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow") + + def model_dump_one_level(self) -> dict[str, Any]: + dumped = {name: getattr(self, name) for name in self.model_fields_set} + dumped.update(self.model_extra or {}) + return dumped + + def _lenient_arg_model(spec: ToolSpec) -> type[ArgModelBase]: - """Build a permissive pydantic arg model for FastMCP's internal validation. + """Build a 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. + whatever the client sent through to ``_dispatch_tool``'s ``**kwargs``, + unknown keys included (``_AllowExtraArgModel`` above) — real + argument-shape enforcement happens once in + ``agent_tools.registry.validate_tool_args``, not duplicated here. + Every declared property is optional/``Any``: each tool already + validates its own required args through that same unified path. """ 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) + return create_model( + f"{spec.name.replace('.', '_')}_Args", __base__=_AllowExtraArgModel, **fields + ) def _make_tool(spec: ToolSpec, *, db_factory: DbFactory) -> Tool: diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index da806746d..9d9f49617 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -65,6 +65,8 @@ def to_mcp_query_response(response: dict[str, Any]) -> dict[str, Any]: """ return { "query": response.get("query"), + "router_used": response.get("router_used"), + "failure_reason": response.get("failure_reason"), "evidence": response.get("evidence") or [], "evidence_text": response.get("evidence_text") or "", "results": response.get("results") or [], @@ -155,8 +157,17 @@ async def query_documents( # automatically before calling run_retrieval_query. # See: shared/services/retrieval/intent/ (to be created) namespace = resolve_mcp_namespace(ctx=ctx) + query = str(query or "").strip() async with db_factory() as db: user_id = await resolve_mcp_user_id(ctx=ctx, db=db) + if not query: + return to_mcp_query_response( + { + "query": query, + "router_used": "empty_query_filtered", + "failure_reason": "empty query — retrieval was not run", + } + ) response = await run_retrieval_query( db=db, user_id=user_id, diff --git a/apps/api/app/services/demo/canonical_bundle.py b/apps/api/app/services/demo/canonical_bundle.py new file mode 100644 index 000000000..fa6923f22 --- /dev/null +++ b/apps/api/app/services/demo/canonical_bundle.py @@ -0,0 +1,259 @@ +"""Publish immutable demo artifacts once per source content version.""" + +from __future__ import annotations + +import asyncio +import json +import re +import shutil +import tempfile +import time +from functools import lru_cache +from hashlib import sha256 +from io import BytesIO +from pathlib import Path +from uuid import uuid4 + +from loguru import logger + +from app.services.demo.canonical_bundle_result import CanonicalDemoBundle +from shared.services.redis.redis_service import RedisService +from shared.services.storage.job_file_storage import JobFileStorage +from shared.services.storage.result_storage import ( + ResultStorage, + UploadedResultBundle, + get_result_storage, +) +from shared.services.storage.storage_adapter import StorageAdapter + +_READY_FILE = ".bundle-ready.json" +_LEASE_SECONDS = 180 +_WAIT_SECONDS = 600 +_RENEW_SCRIPT = """ +if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('expire', KEYS[1], ARGV[2]) +end +return 0 +""" +_RELEASE_SCRIPT = """ +if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('del', KEYS[1]) +end +return 0 +""" +_SourceSignature = tuple[tuple[str, int, int, int], ...] + + +class CanonicalDemoBundleStore: + """Own versioning, upload coordination and the durable completion marker.""" + + def __init__(self, *, redis_service: RedisService) -> None: + self._redis: RedisService = redis_service + self._storage: ResultStorage = get_result_storage() + self._files: JobFileStorage = JobFileStorage() + + async def ensure_bundle( + self, *, source_id: str, source_directory: Path + ) -> CanonicalDemoBundle: + if re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9_-]*", source_id) is None: + raise ValueError("Invalid canonical demo source identifier") + directory: Path = source_directory.resolve(strict=True) + signature: _SourceSignature = await asyncio.to_thread( + _read_source_signature, directory + ) + version: str = await asyncio.to_thread( + _calculate_content_version, directory, signature + ) + storage_id: str = f"demo-canonical/{source_id}/{version}" + bundle: CanonicalDemoBundle | None = await asyncio.to_thread( + self._read_ready_bundle, storage_id, version + ) + if bundle is not None: + return bundle + + lock_key: str = f"lock:demo_bundle:{source_id}:{version}" + owner: str = uuid4().hex + deadline: float = time.monotonic() + _WAIT_SECONDS + while not await self._redis.set_nx(lock_key, owner, ex=_LEASE_SECONDS): + bundle = await asyncio.to_thread( + self._read_ready_bundle, storage_id, version + ) + if bundle is not None: + return bundle + if time.monotonic() >= deadline: + raise TimeoutError("Timed out waiting for canonical demo artifacts") + await asyncio.sleep(0.5) + + stopped: asyncio.Event = asyncio.Event() + renewal: asyncio.Task[None] = asyncio.create_task( + self._renew_lease(lock_key, owner, stopped) + ) + try: + # A previous owner can finish between the first read and SET NX. + bundle = await asyncio.to_thread( + self._read_ready_bundle, storage_id, version + ) + if bundle is not None: + return bundle + upload: asyncio.Task[CanonicalDemoBundle] = asyncio.create_task( + asyncio.to_thread( + self._upload_bundle, storage_id, version, directory, signature + ) + ) + try: + uploaded_bundle: CanonicalDemoBundle = await asyncio.shield(upload) + except asyncio.CancelledError: + # A thread cannot be cancelled; retain the lease until it stops. + await asyncio.gather(upload) + raise + if renewal.done(): + renewal.result() + if not await self._redis.eval( + _RENEW_SCRIPT, keys=[lock_key], args=[owner, _LEASE_SECONDS] + ): + raise RuntimeError("Canonical demo upload lease was lost") + marker_upload: asyncio.Task[None] = asyncio.create_task( + asyncio.to_thread(self._write_ready_marker, uploaded_bundle) + ) + try: + await asyncio.shield(marker_upload) + except asyncio.CancelledError: + await asyncio.gather(marker_upload) + raise + return uploaded_bundle + finally: + stopped.set() + try: + await renewal + except Exception as error: + logger.warning("Canonical demo lease renewal failed: {}", error) + try: + await self._redis.eval(_RELEASE_SCRIPT, keys=[lock_key], args=[owner]) + except Exception as error: + # Ownership expires even if Redis is unavailable during cleanup. + logger.warning("Canonical demo lease release failed: {}", error) + + async def _renew_lease( + self, lock_key: str, owner: str, stopped: asyncio.Event + ) -> None: + while not stopped.is_set(): + try: + await asyncio.wait_for(stopped.wait(), timeout=_LEASE_SECONDS / 3) + except TimeoutError: + if not await self._redis.eval( + _RENEW_SCRIPT, keys=[lock_key], args=[owner, _LEASE_SECONDS] + ): + raise RuntimeError("Canonical demo upload lease was lost") + + def _read_ready_bundle( + self, storage_id: str, version: str + ) -> CanonicalDemoBundle | None: + raw_prefix: str = self._files.build_result_raw_prefix(job_id=storage_id) + marker_key: str = f"{raw_prefix}{_READY_FILE}" + adapter: StorageAdapter = self._files.storage_adapter + bucket: str = self._files.results_bucket + if not adapter.exists(marker_key, bucket): + return None + try: + marker: object = json.loads(adapter.download_fileobj(marker_key, bucket)) + except (ValueError, UnicodeDecodeError): + return None + if not isinstance(marker, dict) or marker.get("content_version") != version: + return None + zip_size: object = marker.get("zip_size") + if type(zip_size) is not int or zip_size <= 0: + return None + zip_key: str = self._files.build_result_zip_key(job_id=storage_id) + if adapter.get_object_size(zip_key, bucket) != zip_size: + return None + return CanonicalDemoBundle( + zip_key=zip_key, + zip_size=zip_size, + raw_prefix=raw_prefix, + content_version=version, + reused=True, + ) + + def _upload_bundle( + self, + storage_id: str, + version: str, + directory: Path, + signature: _SourceSignature, + ) -> CanonicalDemoBundle: + marker_key: str = ( + f"{self._files.build_result_raw_prefix(job_id=storage_id)}{_READY_FILE}" + ) + if self._files.storage_adapter.exists(marker_key, self._files.results_bucket): + # A malformed marker or missing ZIP must stay incomplete during repair. + self._files.delete_object(marker_key, bucket=self._files.results_bucket) + with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temporary: + archive_path: Path = Path( + shutil.make_archive( + str(Path(temporary) / "bundle"), "zip", root_dir=directory + ) + ) + zip_size: int = archive_path.stat().st_size + uploaded: UploadedResultBundle = self._storage.upload( + job_id=storage_id, + result_dir=str(directory), + zip_file_path=str(archive_path), + ) + if _read_source_signature(directory) != signature: + raise RuntimeError("Canonical demo source changed during artifact upload") + return CanonicalDemoBundle( + zip_key=uploaded.zip_key, + zip_size=zip_size, + raw_prefix=uploaded.raw_prefix, + content_version=version, + reused=False, + ) + + def _write_ready_marker(self, bundle: CanonicalDemoBundle) -> None: + # The marker is the commit point; a ZIP alone is never a complete bundle. + payload: bytes = json.dumps( + {"content_version": bundle.content_version, "zip_size": bundle.zip_size} + ).encode("utf-8") + self._files.upload_fileobj( + BytesIO(payload), + f"{bundle.raw_prefix}{_READY_FILE}", + bucket=self._files.results_bucket, + content_type="application/json", + ) + + +def _read_source_signature(directory: Path) -> _SourceSignature: + files: list[tuple[str, int, int, int]] = [] + for path in sorted(directory.rglob("*")): + if path.is_symlink(): + raise ValueError("Canonical demo sources cannot contain symbolic links") + if path.is_file(): + status = path.stat() + files.append( + ( + path.relative_to(directory).as_posix(), + status.st_size, + status.st_mtime_ns, + status.st_ctime_ns, + ) + ) + if not files: + raise ValueError("Canonical demo source directory is empty") + return tuple(files) + + +@lru_cache(maxsize=128) +def _calculate_content_version(directory: Path, signature: _SourceSignature) -> str: + # Stat changes invalidate the local memo; the durable key uses actual bytes. + digest = sha256(b"knowhere-demo-bundle-v1\0") + for relative_path, size, _modified, _changed in signature: + path_bytes: bytes = relative_path.encode("utf-8") + digest.update(len(path_bytes).to_bytes(8, "big")) + digest.update(path_bytes) + digest.update(size.to_bytes(8, "big")) + with (directory / relative_path).open("rb") as source_file: + while block := source_file.read(1024 * 1024): + digest.update(block) + if _read_source_signature(directory) != signature: + raise RuntimeError("Canonical demo source changed while computing its version") + return digest.hexdigest() diff --git a/apps/api/app/services/demo/canonical_bundle_result.py b/apps/api/app/services/demo/canonical_bundle_result.py new file mode 100644 index 000000000..a67fdf142 --- /dev/null +++ b/apps/api/app/services/demo/canonical_bundle_result.py @@ -0,0 +1,12 @@ +"""Immutable object-storage references for one canonical demo version.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CanonicalDemoBundle: + zip_key: str + zip_size: int + raw_prefix: str + content_version: str + reused: bool diff --git a/apps/api/app/services/demo/source_catalog.py b/apps/api/app/services/demo/source_catalog.py index 59c59abea..774234d07 100644 --- a/apps/api/app/services/demo/source_catalog.py +++ b/apps/api/app/services/demo/source_catalog.py @@ -651,9 +651,19 @@ def _normalize_asset_path(asset_path: str) -> Path | None: return Path(*parts) -@lru_cache(maxsize=8) def _load_source_chunks(source: DemoSourceDefinition) -> tuple[dict[str, Any], ...]: chunks_path = (_DATA_ROOT / source.asset_directory) / "chunks.json" + status = chunks_path.stat() + return _load_source_chunks_version( + chunks_path, status.st_size, status.st_mtime_ns, status.st_ctime_ns + ) + + +@lru_cache(maxsize=8) +def _load_source_chunks_version( + chunks_path: Path, size: int, modified_at: int, changed_at: int +) -> tuple[dict[str, Any], ...]: + """Invalidate parsed source content when the canonical input file changes.""" with chunks_path.open("r", encoding="utf-8") as file: payload = json.load(file) diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py index d27f86054..022746046 100644 --- a/apps/api/app/services/demo/source_materializer.py +++ b/apps/api/app/services/demo/source_materializer.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import shutil import tempfile import time @@ -16,6 +17,8 @@ import logfire from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition +from app.services.demo.canonical_bundle import CanonicalDemoBundleStore +from app.services.demo.canonical_bundle_result import CanonicalDemoBundle from sqlalchemy.exc import IntegrityError from sqlalchemy import delete, func, select from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession @@ -33,7 +36,9 @@ from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces from shared.services.retrieval.publication_service import RetrievalPublicationService from shared.services.retrieval.publication_models import DocumentPublicationScope -from shared.services.redis import RedisPublicationSemaphore, RedisServiceFactory +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache +from shared.services.retrieval.scoring.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION +from shared.services.redis import RedisPublicationSemaphore, RedisService, RedisServiceFactory from shared.core.config import settings from shared.services.storage.result_storage import get_result_storage @@ -157,14 +162,18 @@ async def _materialize_source( bind.engine if isinstance(bind, AsyncConnection) else bind ) stage_started_at = time.perf_counter() - result_bundle = _upload_demo_result_bundle( + result_bundle = await _upload_demo_result_bundle( job_id=job_id, + demo_source_id=source.demo_source_id, source_directory=self._catalog.source_directory(source), + redis_service=self._redis_service, ) logfire.info( "Demo materialization source bundle upload completed", demo_source_id=source.demo_source_id, duration_seconds=time.perf_counter() - stage_started_at, + reused=result_bundle.reused, + content_version=result_bundle.content_version, ) db.add( @@ -197,10 +206,15 @@ async def _materialize_source( document_metadata={ "source_file_name": source.title, "demo_source_id": source.demo_source_id, + **( + {"result_raw_prefix": result_bundle.raw_prefix} + if result_bundle.content_version + else {} + ), }, inline_payload={"source": "canonical_demo"}, - result_s3_key=result_bundle["zip_key"], - result_size=result_bundle["zip_size"], + result_s3_key=result_bundle.zip_key, + result_size=result_bundle.zip_size, created_at=timestamp, updated_at=timestamp, ) @@ -259,16 +273,26 @@ async def _materialize_source( demo_source_id=source.demo_source_id, duration_seconds=time.perf_counter() - base_rows_started_at, ) - published_state = await db.run_sync( - lambda sync_db: self._publication_service.publish_document_state( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - chunks=[dict(chunk) for chunk in chunks], - update_namespace_snapshot=False, - trace=trace, + with ( + PublicationPreparationCache.scope( + source_id=source.demo_source_id, + content_version=result_bundle.content_version, + index_format_version=MAP_UNIT_INDEX_FORMAT_VERSION, + ) + if settings.DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED + and result_bundle.content_version + else nullcontext() + ): + published_state = await db.run_sync( + lambda sync_db: self._publication_service.publish_document_state( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + chunks=[dict(chunk) for chunk in chunks], + update_namespace_snapshot=False, + trace=trace, + ) ) - ) await db.run_sync( lambda sync_db: self._publication_service.publish_document_graph( sync_db, @@ -493,32 +517,38 @@ def _materialized_source_payload( ) -def _upload_demo_result_bundle( +async def _upload_demo_result_bundle( *, job_id: str, + demo_source_id: str, source_directory: Path, -) -> dict[str, int | str]: - storage = get_result_storage() - with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temp_directory: - zip_base_path = Path(temp_directory) / job_id - zip_file_path = Path( - shutil.make_archive( - str(zip_base_path), - "zip", - root_dir=source_directory, - ) - ) - zip_size = zip_file_path.stat().st_size - bundle = storage.upload( - job_id=job_id, - result_dir=str(source_directory), - zip_file_path=str(zip_file_path), + redis_service: RedisService, +) -> CanonicalDemoBundle: + if not settings.DEMO_CANONICAL_BUNDLE_ENABLED: + return await asyncio.to_thread( + _upload_job_result_bundle, job_id=job_id, source_directory=source_directory ) + return await CanonicalDemoBundleStore(redis_service=redis_service).ensure_bundle( + source_id=demo_source_id, source_directory=source_directory + ) + - return { - "zip_key": bundle.zip_key, - "zip_size": zip_size, - } +def _upload_job_result_bundle( + *, job_id: str, source_directory: Path +) -> CanonicalDemoBundle: + with tempfile.TemporaryDirectory(prefix="knowhere-demo-result-") as temporary: + archive_path: Path = Path(shutil.make_archive( + str(Path(temporary) / job_id), "zip", root_dir=source_directory + )) + zip_size: int = archive_path.stat().st_size + bundle = get_result_storage().upload( + job_id=job_id, result_dir=str(source_directory), + zip_file_path=str(archive_path), + ) + return CanonicalDemoBundle( + zip_key=bundle.zip_key, zip_size=zip_size, raw_prefix=bundle.raw_prefix, + content_version="", reused=False, + ) def _utc_now() -> datetime: diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index fd97a9ba2..9e1bd839e 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -43,6 +43,7 @@ def _document_chunk_asset_url( *, chunk_type: str, job_id: str | None, + raw_prefix: str | None, file_path: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, @@ -57,8 +58,10 @@ def _document_chunk_asset_url( return None try: - return result_storage.generate_artifact_url( + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=file_path, expires_in=_DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS, ) @@ -71,6 +74,7 @@ def _document_page_assets( *, metadata: dict[str, Any] | None, job_id: str | None, + raw_prefix: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, ) -> list[dict[str, Any]]: @@ -90,6 +94,7 @@ def _document_page_assets( if include_asset_urls and job_id and result_storage is not None: asset_url = _page_asset_url( job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=asset["artifact_ref"], result_storage=result_storage, ) @@ -125,6 +130,7 @@ def _normalize_page_asset(raw_asset: dict[str, Any]) -> dict[str, Any] | None: def _page_asset_url( *, job_id: str, + raw_prefix: str | None, artifact_ref: str, result_storage: ResultStorage, ) -> str | None: @@ -132,8 +138,10 @@ def _page_asset_url( if not normalized_ref or not normalized_ref.startswith("page_citation_assets/"): return None try: - return result_storage.generate_artifact_url( + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=normalized_ref, expires_in=_DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS, ) @@ -150,6 +158,76 @@ def _positive_int(value: Any) -> int | None: return number if number > 0 else None +def _result_raw_prefix(metadata: object) -> str | None: + if not isinstance(metadata, dict): + return None + value = metadata.get("result_raw_prefix") + if not isinstance(value, str): + return None + normalized = value.strip() + return normalized or None + + +def _generate_artifact_url( + storage: ResultStorage, + *, + job_id: str, + raw_prefix: str | None, + artifact_ref: str, + expires_in: int, +) -> str | None: + if raw_prefix is None: + return storage.generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + return storage.generate_artifact_url( + job_id=job_id, + raw_prefix=raw_prefix, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + + +def _verify_raw_exists( + storage: ResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, +) -> bool: + if raw_prefix is None: + return storage.verify_raw_exists(job_id=job_id, relative_path=relative_path) + return storage.verify_raw_exists( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + ) + + +def _generate_raw_file_url( + storage: ResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, + expires_in: int, +) -> str | None: + if raw_prefix is None: + return storage.generate_raw_file_url( + job_id=job_id, + relative_path=relative_path, + expires_in=expires_in, + ) + return storage.generate_raw_file_url( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + expires_in=expires_in, + ) + + def document_payload(document) -> dict[str, Any]: return { "document_id": document.document_id, @@ -266,6 +344,7 @@ async def list_document_chunks( chunk=chunk, section=section, job_id=job_result.job_id, + raw_prefix=_result_raw_prefix(job_result.document_metadata), include_asset_urls=include_asset_urls, result_storage=result_storage, ) @@ -324,6 +403,7 @@ async def get_document_chunk( chunk=chunk, section=section, job_id=job_result.job_id, + raw_prefix=_result_raw_prefix(job_result.document_metadata), include_asset_urls=include_asset_urls, result_storage=result_storage, ), @@ -365,17 +445,26 @@ async def get_document_page_citation_source( return None result_storage = self._result_storage or get_result_storage() - if not result_storage.verify_raw_exists( - job_id=job_result.job_id, - relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, - ): - return None + raw_prefix = _result_raw_prefix(job_result.document_metadata) + try: + if not _verify_raw_exists( + result_storage, + job_id=job_result.job_id, + raw_prefix=raw_prefix, + relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, + ): + return None - source_url = result_storage.generate_raw_file_url( - job_id=job_result.job_id, - relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, - expires_in=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, - ) + source_url = _generate_raw_file_url( + result_storage, + job_id=job_result.job_id, + raw_prefix=raw_prefix, + relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, + expires_in=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, + ) + except Exception as exc: + logger.warning(f"Failed to generate page citation source URL (ignored): {exc}") + return None if not source_url: return None @@ -400,6 +489,7 @@ def _chunk_payload( chunk: DocumentChunk, section: DocumentSection | None, job_id: str | None, + raw_prefix: str | None, include_asset_urls: bool, result_storage: ResultStorage | None, ) -> dict[str, Any]: @@ -409,6 +499,7 @@ def _chunk_payload( page_assets = _document_page_assets( metadata=raw_metadata, job_id=job_id, + raw_prefix=raw_prefix, include_asset_urls=include_asset_urls, result_storage=result_storage, ) @@ -429,6 +520,7 @@ def _chunk_payload( "asset_url": _document_chunk_asset_url( chunk_type=chunk_type, job_id=job_id, + raw_prefix=raw_prefix, file_path=file_path, include_asset_urls=include_asset_urls, result_storage=result_storage, diff --git a/apps/api/scripts/prepare_demo_bundles.py b/apps/api/scripts/prepare_demo_bundles.py new file mode 100644 index 000000000..ef914caaf --- /dev/null +++ b/apps/api/scripts/prepare_demo_bundles.py @@ -0,0 +1,116 @@ +"""Prepare canonical demo result artifacts without materializing documents.""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from pathlib import Path + + +def _bootstrap_python_path() -> None: + api_root: Path = Path(__file__).resolve().parents[1] + shared_root: Path = api_root.parents[1] / "packages" / "shared-python" + for root in (api_root, shared_root): + value: str = os.fspath(root) + if value not in sys.path: + sys.path.insert(0, value) + + +_bootstrap_python_path() + +from app.services.demo.canonical_bundle import CanonicalDemoBundleStore # noqa: E402 +from app.services.demo.source_catalog import ( # noqa: E402 + DemoSourceCatalog, + DemoSourceDefinition, +) +from shared.services.redis.redis_service import RedisService # noqa: E402 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source_ids", metavar="SOURCE_ID", nargs="*") + parser.add_argument( + "--all", action="store_true", help="Select every catalog demo source." + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload missing canonical bundles to the configured result bucket.", + ) + return parser + + +def _select_sources( + parser: argparse.ArgumentParser, + catalog: DemoSourceCatalog, + *, + source_ids: list[str], + all_sources: bool, +) -> list[DemoSourceDefinition]: + if all_sources and source_ids: + parser.error("--all cannot be combined with source IDs") + if not all_sources and not source_ids: + parser.error("provide source IDs or --all") + if all_sources: + return list(catalog.list_sources()) + + sources: list[DemoSourceDefinition] = [] + for source_id in dict.fromkeys(source_ids): + source: DemoSourceDefinition | None = catalog.get_source(source_id) + if source is None: + parser.error(f"unknown demo source ID: {source_id}") + sources.append(source) + return sources + + +async def _upload_sources( + *, catalog: DemoSourceCatalog, sources: list[DemoSourceDefinition] +) -> None: + redis_service: RedisService = RedisService() + try: + store: CanonicalDemoBundleStore = CanonicalDemoBundleStore( + redis_service=redis_service + ) + for source in sources: + bundle = await store.ensure_bundle( + source_id=source.demo_source_id, + source_directory=catalog.source_directory(source), + ) + state: str = "reused" if bundle.reused else "created" + print( + f"{state} {source.demo_source_id} " + f"version={bundle.content_version} " + f"zip={bundle.zip_key} raw_prefix={bundle.raw_prefix}" + ) + finally: + await redis_service.close() + + +def main(argv: list[str] | None = None) -> int: + parser: argparse.ArgumentParser = _build_parser() + arguments: argparse.Namespace = parser.parse_args(argv) + catalog: DemoSourceCatalog = DemoSourceCatalog() + sources: list[DemoSourceDefinition] = _select_sources( + parser, + catalog, + source_ids=arguments.source_ids, + all_sources=arguments.all, + ) + if not arguments.upload: + for source in sources: + print(f"dry-run {source.demo_source_id} {source.title}") + print(f"Selected {len(sources)} source(s); pass --upload to prepare bundles.") + return 0 + + try: + asyncio.run(_upload_sources(catalog=catalog, sources=sources)) + except Exception as error: + print(f"Canonical demo preparation failed: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/api/tests/contract/test_canonical_demo_bundle_contract.py b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py new file mode 100644 index 000000000..443d1647b --- /dev/null +++ b/apps/api/tests/contract/test_canonical_demo_bundle_contract.py @@ -0,0 +1,324 @@ +"""Contract coverage for immutable canonical demo artifact reuse.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from pathlib import Path +from threading import Event +from typing import Literal + +import pytest +from httpx import AsyncClient +from pytest import MonkeyPatch + +from tests.support.contract_database import ContractDatabase + + +def _write_source(directory: Path, content: str) -> None: + directory.mkdir(parents=True, exist_ok=True) + (directory / "source.md").write_text(content, encoding="utf-8") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ("upload", "marker")) +async def test_cancellation_retains_lease_until_storage_work_finishes( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, + monkeypatch: MonkeyPatch, + stage: Literal["upload", "marker"], +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import ( + CanonicalDemoBundleStore, + _calculate_content_version, + _read_source_signature, + ) + from app.services.demo.canonical_bundle_result import CanonicalDemoBundle + from shared.services.redis import RedisServiceFactory + + directory: Path = tmp_path / stage + _write_source(directory, "Cancellation must retain the upload lease") + source_id: str = f"contract-cancel-{stage}" + version: str = _calculate_content_version( + directory, _read_source_signature(directory) + ) + lock_key: str = f"lock:demo_bundle:{source_id}:{version}" + redis_service = RedisServiceFactory.get_service() + store = CanonicalDemoBundleStore(redis_service=redis_service) + loop: asyncio.AbstractEventLoop = asyncio.get_running_loop() + started: asyncio.Event = asyncio.Event() + released: Event = Event() + finished: Event = Event() + original_upload = store._upload_bundle + original_marker = store._write_ready_marker + + def pause_storage() -> None: + loop.call_soon_threadsafe(started.set) + if not released.wait(timeout=10): + raise TimeoutError("Contract did not release paused storage work") + + def upload_bundle( + storage_id: str, + content_version: str, + source_directory: Path, + signature: tuple[tuple[str, int, int, int], ...], + ) -> CanonicalDemoBundle: + pause_storage() + try: + return original_upload( + storage_id, content_version, source_directory, signature + ) + finally: + finished.set() + + def write_marker(bundle: CanonicalDemoBundle) -> None: + pause_storage() + try: + original_marker(bundle) + finally: + finished.set() + + if stage == "upload": + monkeypatch.setattr(store, "_upload_bundle", upload_bundle) + else: + monkeypatch.setattr(store, "_write_ready_marker", write_marker) + request: asyncio.Task[CanonicalDemoBundle] = asyncio.create_task( + store.ensure_bundle(source_id=source_id, source_directory=directory) + ) + try: + await asyncio.wait_for(started.wait(), timeout=10) + owner: object = await redis_service.get(lock_key) + assert owner is not None + request.cancel() + await asyncio.sleep(0) + assert not request.done() + assert not finished.is_set() + assert await redis_service.get(lock_key) == owner + finally: + released.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(request, timeout=10) + + assert finished.is_set() + assert not await redis_service.exists(lock_key) + recovered = await CanonicalDemoBundleStore( + redis_service=redis_service + ).ensure_bundle(source_id=source_id, source_directory=directory) + assert recovered.reused is (stage == "marker") + + +@pytest.mark.asyncio +async def test_canonical_bundle_reuses_and_versions_real_filesystem_objects( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import CanonicalDemoBundleStore + from shared.services.redis import RedisServiceFactory + + source_directory = tmp_path / "source" + _write_source(source_directory, "version one") + store = CanonicalDemoBundleStore( + redis_service=RedisServiceFactory.get_service() + ) + + first = await store.ensure_bundle( + source_id="contract-source", + source_directory=source_directory, + ) + reused = await store.ensure_bundle( + source_id="contract-source", + source_directory=source_directory, + ) + + assert first.reused is False + assert reused.reused is True + assert reused.content_version == first.content_version + assert reused.zip_key == first.zip_key + assert reused.raw_prefix == first.raw_prefix + + _write_source(source_directory, "version two") + changed = await store.ensure_bundle( + source_id="contract-source", + source_directory=source_directory, + ) + + assert changed.reused is False + assert changed.content_version != first.content_version + assert changed.zip_key != first.zip_key + assert changed.raw_prefix != first.raw_prefix + + +@pytest.mark.asyncio +async def test_canonical_bundle_retries_when_completion_marker_is_missing( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import CanonicalDemoBundleStore + from shared.services.redis import RedisServiceFactory + + source_directory = tmp_path / "source" + _write_source(source_directory, "marker recovery") + store = CanonicalDemoBundleStore( + redis_service=RedisServiceFactory.get_service() + ) + initial = await store.ensure_bundle( + source_id="contract-marker-recovery", + source_directory=source_directory, + ) + marker_key = f"{initial.raw_prefix}.bundle-ready.json" + assert store._files.storage_adapter.delete_object( + marker_key, store._files.results_bucket + ) + + original_writer = store._write_ready_marker + + def fail_marker(bundle: object) -> None: + raise RuntimeError("marker write failed") + + store._write_ready_marker = fail_marker # type: ignore[method-assign] + with pytest.raises(RuntimeError, match="marker write failed"): + await store.ensure_bundle( + source_id="contract-marker-recovery", + source_directory=source_directory, + ) + + store._write_ready_marker = original_writer # type: ignore[method-assign] + recovered = await store.ensure_bundle( + source_id="contract-marker-recovery", + source_directory=source_directory, + ) + + assert recovered.reused is False + assert recovered.zip_key == initial.zip_key + assert recovered.content_version == initial.content_version + + +@pytest.mark.asyncio +async def test_canonical_bundle_concurrent_callers_publish_once( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + tmp_path: Path, +) -> None: + async with api_client_factory(): + from app.services.demo.canonical_bundle import CanonicalDemoBundleStore + from shared.services.redis import RedisServiceFactory + + source_directory = tmp_path / "source" + _write_source(source_directory, "concurrent publication") + redis_service = RedisServiceFactory.get_service() + first_store = CanonicalDemoBundleStore(redis_service=redis_service) + second_store = CanonicalDemoBundleStore(redis_service=redis_service) + + results = await asyncio.gather( + first_store.ensure_bundle( + source_id="contract-concurrent", + source_directory=source_directory, + ), + second_store.ensure_bundle( + source_id="contract-concurrent", + source_directory=source_directory, + ), + ) + + assert len({result.zip_key for result in results}) == 1 + assert len({result.raw_prefix for result in results}) == 1 + assert sum(not result.reused for result in results) == 1 + + +@pytest.mark.asyncio +async def test_http_materialization_reuses_canonical_bundle_across_namespaces( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("DEMO_CANONICAL_BUNDLE_ENABLED", "true") + async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + first = await api_client.post( + "/api/v1/demo/materializations", + json={ + "namespace": "canonical-contract-one", + "demo_source_ids": ["demo-tsla-q4-2025"], + }, + ) + second = await api_client.post( + "/api/v1/demo/materializations", + json={ + "namespace": "canonical-contract-two", + "demo_source_ids": ["demo-tsla-q4-2025"], + }, + ) + assert first.status_code == 200 + assert second.status_code == 200 + first_document_id = str(first.json()["sources"][0]["document_id"]) + second_document_id = str(second.json()["sources"][0]["document_id"]) + chunks_response = await api_client.get( + f"/api/v1/documents/{first_document_id}/chunks?page_size=200" + "&include_asset_urls=true" + ) + retrieval_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "canonical-contract-one", + "query": "Tesla investment", + "top_k": 3, + "use_agentic": False, + }, + ) + assert retrieval_response.status_code == 200 + retrieval_asset_urls: list[str] = [ + str(result["asset_url"]) + for result in retrieval_response.json()["results"] + if result.get("asset_url") + ] + assert retrieval_asset_urls + assert all("demo-canonical" in url for url in retrieval_asset_urls) + rows = await ContractDatabase.fetch_all( + """ + SELECT d.document_id, jr.result_s3_key, jr.result_size, + jr.document_metadata ->> 'result_raw_prefix' AS result_raw_prefix + FROM job_results AS jr + JOIN documents AS d ON d.current_job_result_id = jr.id + WHERE d.document_id IN (:first_document_id, :second_document_id) + ORDER BY d.document_id + """, + { + "first_document_id": first_document_id, + "second_document_id": second_document_id, + }, + ) + + monkeypatch.setattr( + source_materializer_module.settings, + "DEMO_CANONICAL_BUNDLE_ENABLED", + False, + ) + existing_document_response = await api_client.get( + f"/api/v1/documents/{first_document_id}/chunks?page_size=200" + "&include_asset_urls=true" + ) + + chunks_body = chunks_response.json() + asset_urls = [ + str(chunk["asset_url"]) + for chunk in chunks_body["chunks"] + if chunk.get("asset_url") + ] + assert asset_urls + assert any("demo-canonical" in asset_url for asset_url in asset_urls) + assert existing_document_response.status_code == 200 + existing_asset_urls = [ + str(chunk["asset_url"]) + for chunk in existing_document_response.json()["chunks"] + if chunk.get("asset_url") + ] + assert existing_asset_urls == asset_urls + assert len(rows) == 2 + assert len({str(row["result_s3_key"]) for row in rows}) == 1 + assert len({str(row["result_raw_prefix"]) for row in rows}) == 1 + assert all(int(row["result_size"]) > 0 for row in rows) diff --git a/apps/api/tests/contract/test_demo_bundle_preparation_contract.py b/apps/api/tests/contract/test_demo_bundle_preparation_contract.py new file mode 100644 index 000000000..d2e8dbcb7 --- /dev/null +++ b/apps/api/tests/contract/test_demo_bundle_preparation_contract.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + + +def _load_script(): + script_path = ( + Path(__file__).resolve().parents[2] / "scripts" / "prepare_demo_bundles.py" + ) + spec = importlib.util.spec_from_file_location("prepare_demo_bundles", script_path) + if spec is None or spec.loader is None: + raise AssertionError("could not load preparation script") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_preparation_defaults_to_explicit_dry_run(capsys) -> None: + module = _load_script() + + assert module.main(["demo-spacex-s1"]) == 0 + output = capsys.readouterr().out + assert "dry-run demo-spacex-s1" in output + assert "--upload" in output + + +def test_preparation_requires_source_selection() -> None: + module = _load_script() + + try: + module.main([]) + except SystemExit as error: + assert error.code == 2 + else: + raise AssertionError("missing source selection should fail") + + +def test_preparation_rejects_all_with_explicit_source(capsys) -> None: + module = _load_script() + + try: + module.main(["--all", "demo-spacex-s1"]) + except SystemExit as error: + assert error.code == 2 + else: + raise AssertionError("mixed source selection should fail") + assert "cannot be combined" in capsys.readouterr().err diff --git a/apps/api/tests/contract/test_demo_preparation_cache_contract.py b/apps/api/tests/contract/test_demo_preparation_cache_contract.py new file mode 100644 index 000000000..20005952a --- /dev/null +++ b/apps/api/tests/contract/test_demo_preparation_cache_contract.py @@ -0,0 +1,167 @@ +"""HTTP and database contract for demo preparation-cache parity.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import Any, cast + +import pytest +from httpx import AsyncClient +from pytest import MonkeyPatch + +from tests.support.contract_database import ContractDatabase + +_SOURCE_ID = "demo-tsla-q4-2025" + + +async def _publish( + client: AsyncClient, + *, + namespace: str, +) -> str: + response = await client.post( + "/api/v1/demo/materializations", + json={"namespace": namespace, "demo_source_ids": [_SOURCE_ID]}, + ) + assert response.status_code == 200 + return str(response.json()["sources"][0]["document_id"]) + + +async def _snapshot(document_id: str) -> dict[str, list[dict[str, Any]]]: + chunks = await ContractDatabase.fetch_all( + """ + SELECT chunk_id, chunk_type, content, content_search_text, + path_search_text, term_search_text, source_chunk_path, + file_path, sort_order + FROM document_chunks + WHERE document_id = :document_id + ORDER BY sort_order, chunk_id, id + """, + {"document_id": document_id}, + ) + units = await ContractDatabase.fetch_all( + """ + SELECT s.section_path, m.unit_kind, m.path_token_count, + m.content_token_count, m.has_image, m.has_table, m.sort_order + FROM document_map_units AS m + JOIN document_sections AS s ON s.section_id = m.section_id + WHERE m.document_id = :document_id + ORDER BY m.sort_order, s.section_path, m.unit_kind, m.id + """, + {"document_id": document_id}, + ) + tokens = await ContractDatabase.fetch_all( + """ + SELECT s.section_path, m.unit_kind, m.sort_order, + t.channel, t.token, t.token_hash, t.frequency + FROM document_map_unit_tokens AS t + JOIN document_map_units AS m ON m.id = t.map_unit_id + JOIN document_sections AS s ON s.section_id = m.section_id + WHERE m.document_id = :document_id + ORDER BY m.sort_order, s.section_path, m.unit_kind, t.channel, t.token + """, + {"document_id": document_id}, + ) + indexes = await ContractDatabase.fetch_all( + """ + SELECT format_version, unit_count, token_count, + average_idf_path, average_idf_content, + path_document_count, path_total_length, + content_document_count, content_total_length + FROM document_map_unit_indexes + WHERE document_id = :document_id + """, + {"document_id": document_id}, + ) + units = sorted(units, key=lambda row: tuple(str(value) for value in row.values())) + tokens = sorted( + tokens, key=lambda row: tuple(str(value) for value in row.values()) + ) + return {"chunks": chunks, "units": units, "tokens": tokens, "indexes": indexes} + + +@pytest.mark.asyncio +async def test_demo_preparation_cache_preserves_persisted_semantics_across_modes( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("DEMO_CANONICAL_BUNDLE_ENABLED", "true") + monkeypatch.setenv("DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", "false") + async with developer_api_client_factory() as client: + from shared.services.retrieval.publication_preparation_cache import ( + PublicationPreparationCache, + ) + + PublicationPreparationCache.clear() + cold_off_id = await _publish(client, namespace="prep-cache-off") + off_snapshot = await _snapshot(cold_off_id) + + import shared.core.config as config_module + + monkeypatch.setattr( + config_module.settings, + "DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", + True, + ) + PublicationPreparationCache.clear() + cache_counts = {"hits": 0, "misses": 0} + original_get = PublicationPreparationCache._get + + def observe_get(key: object) -> object: + value = original_get(key) # type: ignore[arg-type] + cache_counts["hits" if value is not None else "misses"] += 1 + return value + + monkeypatch.setattr( + PublicationPreparationCache, + "_get", + classmethod(lambda _cls, key: observe_get(key)), + ) + cold_on_id = await _publish(client, namespace="prep-cache-cold") + cold_cache_hits = cache_counts["hits"] + cold_cache_misses = cache_counts["misses"] + warm_on_id = await _publish(client, namespace="prep-cache-warm") + assert cold_cache_misses > 0 + assert cache_counts["hits"] > cold_cache_hits + assert cache_counts["misses"] == cold_cache_misses + assert len(PublicationPreparationCache._entries) > 0 + cold_snapshot = await _snapshot(cold_on_id) + warm_snapshot = await _snapshot(warm_on_id) + + monkeypatch.setattr( + config_module.settings, + "DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED", + False, + ) + rollback_id = await _publish(client, namespace="prep-cache-rollback") + rollback_snapshot = await _snapshot(rollback_id) + + for namespace, expected_document_id in ( + ("prep-cache-off", cold_off_id), + ("prep-cache-cold", cold_on_id), + ("prep-cache-warm", warm_on_id), + ("prep-cache-rollback", rollback_id), + ): + retrieval = await client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "Tesla investment", + "top_k": 3, + "use_agentic": False, + }, + ) + assert retrieval.status_code == 200 + result_documents = { + str(result["source"]["document_id"]) + for result in cast(list[dict[str, Any]], retrieval.json()["results"]) + } + assert result_documents + assert result_documents == {expected_document_id} + + assert cold_snapshot == off_snapshot + assert warm_snapshot == off_snapshot + assert rollback_snapshot == off_snapshot diff --git a/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py b/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py index c457e96db..e6d21e389 100644 --- a/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py +++ b/apps/api/tests/contract/test_demo_publication_trace_owner_contract.py @@ -21,13 +21,14 @@ ensure_import_paths() from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition # noqa: E402 +from app.services.demo.canonical_bundle_result import CanonicalDemoBundle # noqa: E402 from app.services.demo.source_materializer import DemoSourceMaterializer # noqa: E402 from shared.models.database.demo_materialization import DemoMaterialization # noqa: E402 from shared.models.database.job import Job # noqa: E402 from shared.services.jobs.lifecycle.publication_trace import PublicationTrace # noqa: E402 from shared.services.retrieval.publication_models import PublishedDocumentState # noqa: E402 from shared.services.retrieval.publication_service import RetrievalPublicationService # noqa: E402 -from shared.services.redis import RedisPublicationSemaphore # noqa: E402 +from shared.services.redis import RedisPublicationSemaphore, RedisService # noqa: E402 class FakeCatalog: @@ -138,6 +139,7 @@ def make_materializer( materializer = object.__new__(materializer_type) materializer._catalog = cast(DemoSourceCatalog, FakeCatalog()) materializer._publication_service = cast(RetrievalPublicationService, publication) + materializer._redis_service = cast(RedisService, object()) return materializer @@ -164,11 +166,16 @@ async def test_demo_owner_finishes_trace_at_publication_commit_or_rollback( updated_at=None, ) - def use_bundle(**kwargs: object) -> dict[str, int | str]: - return { - "zip_key": "zip/123", - "zip_size": 123, - } + async def use_bundle(**kwargs: object) -> CanonicalDemoBundle: + assert kwargs["demo_source_id"] == source.demo_source_id + assert kwargs["source_directory"] == Path("/tmp") + return CanonicalDemoBundle( + zip_key="zip/123", + zip_size=123, + raw_prefix="results/demo-canonical/", + content_version="v1", + reused=False, + ) monkeypatch.setattr(materializer_module, "_upload_demo_result_bundle", use_bundle) monkeypatch.setattr( @@ -179,7 +186,11 @@ def use_bundle(**kwargs: object) -> dict[str, int | str]: monkeypatch.setattr( materializer_module, "settings", - SimpleNamespace(KNOWHERE_PUBLICATION_TRACE_ENABLED=True), + SimpleNamespace( + KNOWHERE_PUBLICATION_TRACE_ENABLED=True, + DEMO_CANONICAL_BUNDLE_ENABLED=True, + DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=False, + ), ) terminal_payloads: list[dict[str, object]] = [] diff --git a/apps/api/tests/contract/test_evidence_renderer_contract.py b/apps/api/tests/contract/test_evidence_renderer_contract.py index da2cac7b4..e6838aee1 100644 --- a/apps/api/tests/contract/test_evidence_renderer_contract.py +++ b/apps/api/tests/contract/test_evidence_renderer_contract.py @@ -51,6 +51,8 @@ def test_mcp_query_response_keeps_evidence_and_debug_results() -> None: assert response == { "query": "q", + "router_used": None, + "failure_reason": None, "evidence": [{"type": "text", "text": "t"}], "evidence_text": "t", "results": [ diff --git a/apps/api/tests/contract/test_publication_chunk_copy_contract.py b/apps/api/tests/contract/test_publication_chunk_copy_contract.py index ce9bf36e3..57a3383ef 100644 --- a/apps/api/tests/contract/test_publication_chunk_copy_contract.py +++ b/apps/api/tests/contract/test_publication_chunk_copy_contract.py @@ -6,6 +6,8 @@ from collections.abc import Iterator import pytest +from psycogreen.gevent import patch_psycopg +from psycopg2.extensions import get_wait_callback, set_wait_callback from sqlalchemy import create_engine, text from sqlalchemy.engine import make_url from sqlalchemy.ext.asyncio import create_async_engine @@ -111,9 +113,17 @@ def _read_rows(connection) -> list[tuple[object, ...]]: ) +@pytest.mark.parametrize("has_cooperative_callback", [False, True]) def test_candidate_chunk_copy_preserves_null_empty_json_and_rollback( contract_database_url: str, + has_cooperative_callback: bool, ) -> None: + original_callback = get_wait_callback() + if has_cooperative_callback: + patch_psycopg() + else: + set_wait_callback(None) + expected_callback = get_wait_callback() url = make_url(contract_database_url).set(drivername="postgresql+psycopg2") engine = create_engine(url) try: @@ -122,6 +132,7 @@ def test_candidate_chunk_copy_preserves_null_empty_json_and_rollback( connection.commit() with Session(bind=connection) as session: insert_chunks_with_copy(session, _build_chunks()) + assert get_wait_callback() is expected_callback rows = _read_rows(connection) assert rows[0][1:5] == ( "", @@ -135,6 +146,7 @@ def test_candidate_chunk_copy_preserves_null_empty_json_and_rollback( assert connection.scalar(text("SELECT count(*) FROM document_chunks")) == 0 finally: engine.dispose() + set_wait_callback(original_callback) async def test_candidate_chunk_copy_supports_asyncpg_owner_and_rollback( diff --git a/apps/api/tests/contract/test_publication_token_copy_contract.py b/apps/api/tests/contract/test_publication_token_copy_contract.py index 8f651c1a4..e2e869ff1 100644 --- a/apps/api/tests/contract/test_publication_token_copy_contract.py +++ b/apps/api/tests/contract/test_publication_token_copy_contract.py @@ -6,8 +6,10 @@ from collections.abc import Iterator import pytest +from psycogreen.gevent import patch_psycopg +from psycopg2.extensions import get_wait_callback, set_wait_callback from sqlalchemy import create_engine, text -from sqlalchemy.engine import make_url +from sqlalchemy.engine import Connection, make_url from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.orm import Session @@ -59,9 +61,17 @@ def contract_database_url( yield get_contract_database_url() +@pytest.mark.parametrize("has_cooperative_callback", [False, True]) def test_candidate_copy_preserves_rows_and_rollback_with_psycopg2( contract_database_url: str, + has_cooperative_callback: bool, ) -> None: + original_callback = get_wait_callback() + if has_cooperative_callback: + patch_psycopg() + else: + set_wait_callback(None) + expected_callback = get_wait_callback() url = make_url(contract_database_url).set(drivername="postgresql+psycopg2") engine = create_engine(url) try: @@ -70,6 +80,7 @@ def test_candidate_copy_preserves_rows_and_rollback_with_psycopg2( connection.commit() with Session(bind=connection) as session: insert_token_rows_with_copy(session, _TOKEN_ROWS) + assert get_wait_callback() is expected_callback values = connection.execute( text("SELECT token, frequency FROM document_map_unit_tokens ORDER BY id") ).all() @@ -80,6 +91,7 @@ def test_candidate_copy_preserves_rows_and_rollback_with_psycopg2( ) == 0 finally: engine.dispose() + set_wait_callback(original_callback) async def test_candidate_copy_preserves_rows_and_rollback_with_asyncpg( @@ -114,3 +126,57 @@ def copy_and_read(sync_connection) -> list[tuple[str, int]]: ) == 0 finally: await engine.dispose() + + +async def test_candidate_token_writes_complete_under_per_command_budget( + contract_database_url: str, +) -> None: + """Slow token writes must stay atomic without requiring one long command.""" + url = make_url(contract_database_url).set(drivername="postgresql+asyncpg") + engine = create_async_engine( + url, + connect_args={ + "command_timeout": 3, + "server_settings": {"statement_timeout": "3000"}, + }, + ) + rows: list[dict[str, object]] = [ + {**_TOKEN_ROWS[0], "id": f"dmut_slow_{index}"} + for index in range(2_501) + ] + try: + async with engine.connect() as connection: + await connection.execute(text(_CREATE_TOKEN_TABLE)) + await connection.execute(text(""" + CREATE FUNCTION pg_temp.delay_token_write() RETURNS trigger + LANGUAGE plpgsql AS $$ + BEGIN + PERFORM pg_sleep((SELECT count(*) FROM inserted_tokens) * 0.002); + RETURN NULL; + END $$ + """)) + await connection.execute(text(""" + CREATE TRIGGER delay_token_write AFTER INSERT + ON document_map_unit_tokens + REFERENCING NEW TABLE AS inserted_tokens + FOR EACH STATEMENT EXECUTE FUNCTION pg_temp.delay_token_write() + """)) + await connection.commit() + transaction = await connection.begin() + # Publication has already written its document before token persistence. + await connection.execute(text("SELECT 1")) + + def write_tokens(sync_connection: Connection) -> None: + with Session(bind=sync_connection) as session: + insert_token_rows_with_copy(session, rows) + + await connection.run_sync(write_tokens) + assert await connection.scalar( + text("SELECT count(*) FROM document_map_unit_tokens") + ) == len(rows) + await transaction.rollback() + assert await connection.scalar( + text("SELECT count(*) FROM document_map_unit_tokens") + ) == 0 + finally: + await engine.dispose() diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index cffa73b38..708c13921 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -303,58 +303,6 @@ async def test_retrieval_should_use_classic_topk_when_agentic_is_false( assert response_json["router_used"] == "classic_topk" -async def test_should_return_request_validation_failure_for_an_invalid_channel( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - async with developer_api_client_factory() as api_client: - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "default", - "query": "alpha", - "channels": ["invalid-channel"], - }, - ) - - assert response.status_code == 400 - assert response.headers["x-request-id"] - - response_json = cast(dict[str, object], response.json()) - error = cast(dict[str, object], response_json["error"]) - details = cast(dict[str, object], error["details"]) - violations = cast(list[dict[str, object]], details["violations"]) - - assert response_json["success"] is False - assert error["code"] == "INVALID_ARGUMENT" - assert error["message"] == "Request validation failed" - assert violations[0]["field"] == "body.channels" - assert "Invalid channel" in cast(str, violations[0]["description"]) - - -async def test_should_reject_legacy_channel_controls( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - async with developer_api_client_factory() as api_client: - response = await api_client.post( - "/api/v1/retrieval/query", - json={ - "namespace": "default", - "query": "alpha", - "channels": ["content"], - }, - ) - - assert response.status_code == 400 - response_json = cast(dict[str, object], response.json()) - error = cast(dict[str, object], response_json["error"]) - assert error["code"] == "INVALID_ARGUMENT" - assert "deprecated and unsupported" in str(error) - - async def test_should_exclude_matching_document_ids_from_the_response( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] @@ -935,8 +883,6 @@ async def complete(self, *_args: object, **_kwargs: object) -> None: chunk_types=None, signal_paths=None, filter_mode="delete", - channels=None, - channel_weights=None, rerank=False, threshold=0.0, internal_recall_k=None, diff --git a/apps/api/tests/contract/test_retrieval_document_scope.py b/apps/api/tests/contract/test_retrieval_document_scope.py index 9329414df..51fbc7e8a 100644 --- a/apps/api/tests/contract/test_retrieval_document_scope.py +++ b/apps/api/tests/contract/test_retrieval_document_scope.py @@ -142,20 +142,7 @@ async def call(name, args): for name, args in [ ("list_documents", {}), ("grep", {"pattern": "scopeprobe"}), - ("grep", {"pattern": "scopeprobe", "document_ids": ids}), - ("recall", {"query": "scopeprobe", "channels": ["term"], "top_k": 50}), - ( - "recall", - {"query": "scopeprobe", "channels": ["path_content"], "top_k": 50}, - ), - ("recall", {"query": "scopeprobe", "document_ids": ids, "top_k": 50}), - ( - "node_filter", - { - "document_ids": ids, - "predicates": [{"field": "path", "terms": ["Section"]}], - }, - ), + ("recall", {"query": "scopeprobe", "limit": 50}), ("assets", {}), ( "assets", @@ -177,21 +164,53 @@ async def call(name, args): ) if expected: assert result.error is None, (name, result.error) + # scope only ever names documents this call is actually allowed + # to see (an out-of-boundary document_id fails resolve_scope, by + # design — see agent_tools/scope.py) — a legitimately-scoped + # call must still narrow correctly to those docs. + if expected: + scope_expected = [{"document_id": doc_id} for doc_id in sorted(expected)] + for name, args in [ + ("grep", {"pattern": "scopeprobe", "scope": scope_expected}), + ( + "recall", + {"query": "scopeprobe", "scope": scope_expected, "limit": 50}, + ), + ( + "node_filter", + { + "scope": scope_expected, + "predicates": [{"field": "path", "terms": ["Section"]}], + }, + ), + ("assets", {"scope": scope_expected}), + ]: + result = await call(name, args) + assert result.error is None, (name, result.error) + assert {ref["document_id"] for ref in result.refs} == expected, ( + name, + include, + exclude, + result, + ) for doc_id in ids: - outline = await call("outline", {"document_id": doc_id}) + outline = await call("outline", {"scope": [{"document_id": doc_id}]}) assert {r["document_id"] for r in outline.refs} == ({doc_id} & expected) neighbors = await call("neighbors", {"document_id": doc_id}) assert {r["document_id"] for r in neighbors.refs} == ( expected - {doc_id} if doc_id in expected else set() ) - # Explicit tool selectors cannot widen the request boundary. + # Explicit tool selectors cannot widen the request boundary: naming + # an out-of-boundary document_id in scope fails resolve_scope, so + # no refs (and no data) ever come back for it. if expected: outside = next((doc for doc in ids if doc not in expected), None) if outside: + outside_scope = [{"document_id": outside}] for name, args in [ - ("grep", {"pattern": "scopeprobe", "document_ids": [outside]}), - ("recall", {"query": "scopeprobe", "document_ids": [outside]}), - ("assets", {"document_ids": [outside]}), + ("grep", {"pattern": "scopeprobe", "scope": outside_scope}), + ("recall", {"query": "scopeprobe", "scope": outside_scope}), + ("assets", {"scope": outside_scope}), ]: assert not (await call(name, args)).refs async with contract_db_session() as db: @@ -528,9 +547,13 @@ def verify_observation(content): id="grep", function=SimpleNamespace( name="corpus_grep", - arguments=json.dumps( - {"pattern": "scopeprobe", "document_ids": ids} - ), + # No explicit scope: document_scope (already bound into + # this episode) is the only boundary in effect here — + # naming an out-of-boundary document_id in scope would + # fail the call instead of silently narrowing (see + # agent_tools/scope.py), which isn't what this test + # exercises. + arguments=json.dumps({"pattern": "scopeprobe"}), ), ), ] diff --git a/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py b/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py deleted file mode 100644 index e056af6cc..000000000 --- a/apps/api/tests/contract/test_retrieval_rrf_duplicate_contract.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Contracts for duplicate chunk handling in reciprocal-rank fusion.""" - -from __future__ import annotations - -from shared.services.retrieval.search.scoring import merge_channels_rrf - - -def test_rrf_counts_each_chunk_once_per_channel() -> None: - rows = [ - {"chunk_id": "shared", "document_id": "doc-a"}, - {"chunk_id": "shared", "document_id": "doc-b"}, - {"chunk_id": "other", "document_id": "doc-c"}, - ] - - result = merge_channels_rrf([rows], [1.0], top_k=3) - - assert [row["chunk_id"] for row in result] == ["shared", "other"] - assert result[0]["score"] == round(1.0 / 61, 6) - assert result[1]["score"] == round(1.0 / 62, 6) diff --git a/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py b/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py new file mode 100644 index 000000000..f3ff96473 --- /dev/null +++ b/apps/api/tests/migrations/test_legacy_token_lookup_retirement_contract.py @@ -0,0 +1,148 @@ +"""Real PostgreSQL contracts for gated and replayable token index maintenance.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +from alembic import command +from alembic.config import Config +from sqlalchemy import text +from sqlalchemy.engine import Engine + +from shared.services.retrieval.maintenance.legacy_token_lookup import ( + LegacyTokenLookupMaintenance, +) + +_ROOT: Path = Path(__file__).resolve().parents[4] +_PARENT_REVISION: str = "2b3c4d5e6f7a" +_REVISION: str = "3c4d5e6f7a8b" +_LEGACY_INDEX: str = "idx_document_map_unit_tokens_lookup" +_BINARY_INDEX: str = "idx_document_map_unit_tokens_token_lookup_binary" +_UNIT_LOOKUP_INDEX: str = "idx_document_map_unit_tokens_unit_lookup" + + +def _build_config() -> Config: + config: Config = Config(str(_ROOT / "apps/api/alembic.ini")) + config.set_main_option("script_location", str(_ROOT / "apps/api/alembic")) + return config + + +def _inspect_indexes(engine: Engine) -> dict[str, dict[str, str | bool]]: + with engine.connect() as connection: + return LegacyTokenLookupMaintenance(connection).inspect_indexes() + + +def _run_operator(*arguments: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(_ROOT / "scripts/retire_legacy_token_lookup.py"), *arguments], + env=os.environ.copy(), capture_output=True, text=True, timeout=30, check=False, + ) + + +def test_should_keep_default_indexes_then_apply_and_restore_at_same_head( + alembic_engine: Engine, +) -> None: + config: Config = _build_config() + command.upgrade(config, "heads") + initial: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + assert len(initial) == 5 + assert _LEGACY_INDEX in initial + inspection: subprocess.CompletedProcess[str] = _run_operator() + assert inspection.returncode == 0, inspection.stderr + assert json.loads(inspection.stdout) == initial + + for _ in range(2): + result: subprocess.CompletedProcess[str] = _run_operator("--apply") + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == { + name: state for name, state in initial.items() if name != _LEGACY_INDEX + } + with alembic_engine.connect() as connection: + assert connection.scalar(text("SELECT version_num FROM alembic_version")) == _REVISION + for _ in range(2): + result = _run_operator("--restore") + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == initial + + +@pytest.mark.parametrize("is_external_transaction", (False, True)) +def test_should_gate_upgrade_and_restore_on_downgrade( + alembic_engine: Engine, is_external_transaction: bool, +) -> None: + config: Config = _build_config() + command.upgrade(config, _PARENT_REVISION) + initial: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + config.cmd_opts = argparse.Namespace(x=["retire_legacy_token_lookup=true"]) + if is_external_transaction: + with alembic_engine.begin() as connection: + config.attributes["connection"] = connection + command.upgrade(config, _REVISION) + config.attributes.pop("connection") + else: + command.upgrade(config, _REVISION) + assert _inspect_indexes(alembic_engine) == { + name: state for name, state in initial.items() if name != _LEGACY_INDEX + } + if is_external_transaction: + with alembic_engine.begin() as connection: + config.attributes["connection"] = connection + command.downgrade(config, _PARENT_REVISION) + else: + command.downgrade(config, _PARENT_REVISION) + assert _inspect_indexes(alembic_engine) == initial + + +@pytest.mark.parametrize( + ("index_name", "damage"), + ( + (_BINARY_INDEX, "missing"), + (_UNIT_LOOKUP_INDEX, "incompatible"), + (_BINARY_INDEX, "invalid"), + (_UNIT_LOOKUP_INDEX, "not_ready"), + ), +) +def test_should_refuse_retirement_without_compatible_replacements( + alembic_engine: Engine, index_name: str, damage: str, +) -> None: + command.upgrade(_build_config(), "heads") + with alembic_engine.begin() as connection: + if damage in {"missing", "incompatible"}: + connection.execute(text(f"DROP INDEX {index_name}")) + if damage == "incompatible": + connection.execute(text( + f"CREATE INDEX {index_name} ON document_map_unit_tokens (channel)" + )) + else: + column: str = "indisvalid" if damage == "invalid" else "indisready" + connection.execute(text( + f"UPDATE pg_index SET {column} = false " + "WHERE indexrelid = to_regclass(:name)" + ), {"name": index_name}) + before: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + result: subprocess.CompletedProcess[str] = _run_operator("--apply") + assert result.returncode == 1 + assert index_name in result.stderr + assert _inspect_indexes(alembic_engine) == before + + +def test_should_refuse_to_replace_incompatible_legacy_index( + alembic_engine: Engine, +) -> None: + command.upgrade(_build_config(), "heads") + with alembic_engine.begin() as connection: + connection.execute(text(f"DROP INDEX {_LEGACY_INDEX}")) + connection.execute(text( + f"CREATE INDEX {_LEGACY_INDEX} ON document_map_unit_tokens (channel)" + )) + before: dict[str, dict[str, str | bool]] = _inspect_indexes(alembic_engine) + for argument in ("--restore", "--apply"): + result: subprocess.CompletedProcess[str] = _run_operator(argument) + assert result.returncode == 1 + assert _LEGACY_INDEX in result.stderr + assert _inspect_indexes(alembic_engine) == before diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index a5a64051a..ff6ed485f 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -270,6 +270,37 @@ def test_should_index_document_chunks_in_lazy_section_order( ) +def test_should_not_have_map_unit_term_search_text( + migrated_head_engine: Engine, +) -> None: + with migrated_head_engine.begin() as connection: + columns = connection.execute( + text( + """ + SELECT column_name + FROM information_schema.columns + WHERE table_schema = current_schema() + AND table_name = 'document_map_units' + AND column_name = 'term_search_text_lower' + """ + ) + ).fetchall() + index_count = connection.execute( + text( + """ + SELECT count(*) + FROM pg_class + JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace + WHERE pg_namespace.nspname = current_schema() + AND pg_class.relname = 'idx_document_map_units_term_trgm' + """ + ) + ).scalar_one() + + assert columns == [] + assert index_count == 0 + + def test_should_not_have_content_trigram_index( migrated_head_engine: Engine, ) -> None: diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py index bfebab1d1..f42f219b7 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -267,6 +267,25 @@ def test_result_storage_allows_page_citation_and_page_pdf_artifact_refs_not_debu ) +def test_result_storage_uses_validated_canonical_raw_prefix_without_job_fallback() -> None: + storage = JobResultStorage(results_bucket="test-results") + + assert ( + storage.build_raw_key( + job_id="job-materialization", + raw_prefix="results/demo-canonical/source-a/content-v1/", + relative_path="images/chart.png", + ) + == "results/demo-canonical/source-a/content-v1/images/chart.png" + ) + with pytest.raises(ValueError, match="under results"): + storage.build_raw_key( + job_id="job-materialization", + raw_prefix="uploads/shared/", + relative_path="images/chart.png", + ) + + def test_result_storage_upload_filters_to_referenced_artifacts(tmp_path) -> None: class FakeStorageAdapter: def __init__(self) -> None: diff --git a/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py b/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py index f5ca83f6b..bf72b95a9 100644 --- a/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py +++ b/deprecated/mapnav/tests/api-contract/test_retrieval_mapnav_session_contract.py @@ -204,8 +204,6 @@ async def fake_assemble_retrieval_results( chunk_types=None, signal_paths=None, filter_mode="delete", - channels=None, - channel_weights=None, rerank=False, threshold=0.0, internal_recall_k=None, diff --git a/docs/design/canonical-demo-bundle-reuse.md b/docs/design/canonical-demo-bundle-reuse.md new file mode 100644 index 000000000..98e3d8336 --- /dev/null +++ b/docs/design/canonical-demo-bundle-reuse.md @@ -0,0 +1,136 @@ +# Canonical demo bundle reuse + +## Scope + +Demo materialization reuses the same immutable ZIP and raw artifacts across +users and namespaces. Jobs, job results, documents, publication transactions, +authorization and retrieval isolation remain owner scoped. Normal Worker parsing +continues to store its result artifacts by job identifier. + +The previous demo path compressed and uploaded every artifact for each new job. +Two real local SpaceX materializations created 637 objects each and took 40.422 s +and 39.098 s from HTTP request to response. These are local end-to-end timings, +not production measurements or isolated publication transaction timings. +Each request duplicated 439,190,627 stored bytes, including a 215,559,924-byte ZIP. + +After the change, the first canonical upload took 40.211 s. Two subsequent +materializations into separate namespaces took **5.483 s** and **5.337 s**, +respectively (about 86% lower than the two legacy samples). Both added zero +objects and rewrote zero objects; all 638 canonical objects retained their +ETags, sizes and modification times. The extra object is the completion marker. +This removes about 439 MB of duplicate object writes per repeat materialization. + +Both repeat requests passed real HTTP checks of document image/table/page +assets, classic retrieval page/image assets, the public job ZIP URL, and archive +readback. Archiving one materialization left the other's shared assets readable. + +## Storage contract + +The content version hashes sorted relative file paths, lengths and file bytes, +including a bundle-format version. A bounded process-local memo avoids rereading +unchanged file bytes; changes to a file's size, mtime or ctime invalidate it. +Symbolic links and sources that change during hashing/upload are rejected. + +Keys preserve the existing result storage layout: + +```text +results/demo-canonical//.zip +results/demo-canonical/// +results/demo-canonical///.bundle-ready.json +``` + +The marker is written only after ZIP and raw uploads complete. A failed partial +upload is never reused. A subsequent request can retry it. A renewable Redis +lease serializes first uploads across API processes; release compares the owner +token and an abandoned lease expires. Compression, hashing and storage calls run +off the API event loop. Cancellation waits for an active upload thread before +releasing its lease. + +Reuse checks the completion marker and ZIP size without issuing one storage +request per raw artifact. If a raw object is deleted after completion, reuse +does not detect or repair it automatically. Recovery requires invalidating that +version's `.bundle-ready.json` and rerunning the preparation command against the +same source version. Protect the shared prefix from per-job lifecycle cleanup. + +`JobResult.result_s3_key` records the canonical ZIP. +`JobResult.document_metadata.result_raw_prefix` records the canonical raw prefix. +Artifact readers use that metadata when present and retain the job-key fallback +for older results and ordinary Worker output. No database migration is required. + +## DevOps enablement and rollback + +1. Deploy the writer and compatible artifact readers together. The independent + `DEMO_CANONICAL_BUNDLE_ENABLED` setting defaults to `false`. +2. Prepare bundles with the same deployed source files and result bucket, using + the [preparation command](canonical-demo-preparation.md). This moves the + initial ZIP compression and artifact upload ahead of user requests. A new + unprepared content version still needs its first upload. +3. Enable `DEMO_CANONICAL_BUNDLE_ENABLED=true` on API instances after checking + access to their configured result bucket and Redis. This does not change + `KNOWHERE_PUBLICATION_STRATEGY` or retire any database index. +4. Materialize the same source into two namespaces. Verify identical ZIP and + raw keys, and unchanged object timestamps/ETags on the second request. Fetch + image/table/page artifacts and the ZIP through their public response URLs. +5. Roll back new writes by setting `DEMO_CANONICAL_BUNDLE_ENABLED=false` and + restarting the API. Keep the compatible readers deployed: they continue to + resolve existing canonical results while new jobs use individual bundles. +6. Do not delete canonical objects on per-document archive or expiration. A + source/version can have references from multiple owners. Bucket lifecycle + rules must preserve this prefix for the full lifetime of its references. + +A binary rollback to a release without prefix-aware artifact readers requires +first restoring per-job raw objects for every canonical result, or retaining +the reader compatibility patch. Turning the switch off does not rewrite +existing result records. + +The separate `DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED` flag also defaults to +`false`. It requires canonical bundles and reuses only pure lexical preparation +inside the current API process. See [cache boundaries and measurements](demo-publication-preparation-cache.md). + +[Legacy token lookup retirement](legacy-token-lookup-retirement.md) is an +independent database operation. Normal deployment preserves that index; enabling +either demo flag does not change it. Index restoration uses the operator command +at the same Alembic head and can take substantially longer than retirement. + +## Verification boundary + +Contract coverage targets reuse across namespaces, content changes, failed +uploads, concurrent creation and legacy artifact compatibility. Runtime +verification must also exercise actual HTTP, Redis/Celery Worker execution, +real parsing, publication, retrieval and archive. Local results do not establish +production latency or performance at a larger database size. + +The real HTTP verifier caught a missing raw-prefix projection in classic +`map_unit_discovery` hydration: materialization and document asset endpoints +succeeded, while a retrieved page image returned 404. The query now projects +the artifact prefix alongside the job identifier. Artifact location is also +carried through connected/reference hydration, small-corpus loading and the +agent read/grep/table tools; their storage reads share the same prefix support. + +The actual Worker Markdown regression used presigned local S3 upload and +confirm-upload, followed by Redis/Celery dispatch, the real parser, publication, +classic retrieval and archive. It completed in 3.240 s observed by the client; +the parser trace was 2.307 s. This small document verifies the ordinary Worker +path and is not a SpaceX parse benchmark. +After the canonical-reader changes, the Worker was restarted and the same real +queued path passed again in 3.278 s (job duration 2.906 s). All runtime test +documents were archived. Production configuration and indexes were untouched. + +After the preparation-cache changes, a fresh API and Worker processed another +HTTP-created Markdown job in 4.300 s client-observed time (3.814 s job duration). +The fresh Worker log confirmed ownership; two chunks, classic retrieval and +archive passed. The configured summarization provider returned HTTP 401, so this +run exercised the parser's fallback. It does not validate successful external +LLM/PDF parsing or full SpaceX parsing latency. + +Focused verification passed: six canonical/trace contracts, eight storage and +document lifecycle contracts, and 48 existing retrieval checks. The canonical +HTTP contract also exercises classic retrieval's shared asset URL after the +raw-SQL projection fix. Ruff, targeted Pyright and `git diff --check` passed. + +Run the canonical regression contracts with: + +```sh +uv run pytest -q apps/api/tests/contract/test_canonical_demo_bundle_contract.py \ + apps/api/tests/contract/test_demo_publication_trace_owner_contract.py +``` diff --git a/docs/design/canonical-demo-preparation.md b/docs/design/canonical-demo-preparation.md new file mode 100644 index 000000000..a97d8d827 --- /dev/null +++ b/docs/design/canonical-demo-preparation.md @@ -0,0 +1,27 @@ +# Prepare canonical demo bundles + +`prepare_demo_bundles.py` uploads the immutable ZIP and raw artifacts for +catalog demo sources ahead of user materialization. It uses the same content +version, Redis lease and completion marker as the API path. It does not create +jobs, documents, credits or publication state. + +Run from the repository root with the API's normal storage and Redis environment +configured (`S3_TYPE`, result bucket, endpoint and credentials, `REDIS_HOST` and +`REDIS_PORT`). The source checkout must include `apps/api/app/data/demo_documents`. + +```sh +uv run --package knowhere-api-app python apps/api/scripts/prepare_demo_bundles.py demo-spacex-s1 +uv run --package knowhere-api-app python apps/api/scripts/prepare_demo_bundles.py demo-spacex-s1 --upload +uv run --package knowhere-api-app python apps/api/scripts/prepare_demo_bundles.py --all --upload +``` + +The default invocation lists selected sources and makes no storage or Redis +calls. `--upload` is required for writes. A completed version is reported as +`reused` and uploads no objects. New versions are reported as `created` with +their exact ZIP key and raw prefix. A failure exits nonzero; rerunning is safe +because the completion marker is written only after all artifacts finish. + +Use the same configured result bucket as the API. The +`DEMO_CANONICAL_BUNDLE_ENABLED` API setting controls new materializations, not +this explicit preparation command. Keep the prefix-aware readers deployed for +existing canonical results. diff --git a/docs/design/demo-publication-preparation-cache.md b/docs/design/demo-publication-preparation-cache.md new file mode 100644 index 000000000..6c9039194 --- /dev/null +++ b/docs/design/demo-publication-preparation-cache.md @@ -0,0 +1,61 @@ +# Demo publication preparation cache + +## Purpose + +Canonical demo artifacts can be reused across user namespaces. Publishing each +copy still tokenizes the same parsed text and builds the same token-frequency +counts. A bounded process-local cache can reuse those pure results without +changing publication ownership or retrieval semantics. + +## Boundary + +The cache is enabled only when `DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED=true` +and the canonical bundle supplies a content version. Its scope identifies the +source, content version, map-unit index format, and preparation algorithm +version. Other publication paths continue through the authoritative tokenizers. + +Cache values are immutable token sequences or frequency pairs. Callers receive +fresh mutable copies. Document IDs, job result IDs, section IDs, map-unit IDs, +namespace, transaction state, scores, and retrieval results never enter the +cache. Search-field construction and map-unit persistence remain in the shared +publication path and retain the same ordering and transaction boundary. + +The process cache is limited to 64 MiB and 4,096 entries. Eviction only causes +recomputation. A changed source content version, index format, or preparation +algorithm does not reuse an earlier entry. The default flag is off, so rollback +is a configuration change and existing persisted documents remain readable. + +## Verification + +The contract publishes the same canonical demo source in separate namespaces +with the flag off, on with an empty cache, on with a warm cache, and off again. +It compares persisted chunk search fields, map-unit semantics, token-frequency +rows, and IDF statistics after excluding generated database identifiers. It +also checks namespace-scoped retrieval and confirms no document crosses scope. + +This verifies semantic parity. Performance admission still requires a measured +publication stage improvement and end-to-end timing on production-shaped input. + +## Focused measurements + +The catalog's real SpaceX demo (227 chunks) was exercised through the real HTTP +materializer against the small local database, with the canonical bundle +already present. The separate 12.42-million-token clone experiment was not +repeated for this cache: + +| Mode | Request 1 | Request 2 | Bundle writes | +| --- | ---: | ---: | --- | +| Preparation cache off | 4.887 s | 5.051 s | 0 new or changed objects | +| Preparation cache on | 7.044 s (cold) | 2.967 s (warm) | 0 new or changed objects | + +Both modes used the same canonical ZIP and raw prefix. Page and image assets, +classic retrieval, and archival completed successfully. The warm sample is an +observed local improvement; the cold sample includes cache population and was +slower. These are focused observations, not a production p95 claim. + +The final Worker owner check used a small Markdown upload because the local +provider credentials do not establish successful external LLM or PDF parsing. +The Worker completed job `job_f5d5c6d4ed2b` in 3.814 s with a 236 ms publication +trace, and retrieval plus archival passed. The provider returned HTTP 401 and +the parser completed through its fallback; this validates the Worker publication boundary and +fallback behavior, not external model quality. diff --git a/docs/design/legacy-token-lookup-retirement.md b/docs/design/legacy-token-lookup-retirement.md new file mode 100644 index 000000000..b7ef66124 --- /dev/null +++ b/docs/design/legacy-token-lookup-retirement.md @@ -0,0 +1,122 @@ +# Legacy token lookup retirement + +## Purpose and evidence + +Retire only the nonunique `idx_document_map_unit_tokens_lookup` index to reduce +token persistence I/O. Keep the primary key, binary token lookup, unit index, +and unit-leading covering lookup. + +The production incident's token COPY samples waited on `IO:DataFileRead`. +In the disposable 12.42-million-row clone, removing this legacy index reduced +COPY buffer reads by approximately 53%, WAL by 47%, database COPY execution time +by 48%, and application token-stage time by 38%. Publication improved from +11.376 to 10.182 seconds in that pair; full HTTP time stayed around 54 seconds. +These are three targeted cold-cache publication observations, not production +latency guarantees. See [the experiment report](publication-scale-cold-cache-findings.md). + +Current classic discovery and recall use binary lookup. The remaining plain +token-hash reader in `scoring/map_lighting.py` also scopes by map-unit IDs and +channel, which the retained unit-leading index covers. Ten frozen query cases +returned identical counts and hashes before and after retirement; each used +the retained unit lookup. Historical application versions and every possible +query shape have not been validated. + +## Deployment behavior + +Revision `3c4d5e6f7a8b`, following `2b3c4d5e6f7a`, retains the legacy index by +default. No runtime application flag deletes or recreates this index. The +previous retirement migration `2f3g4h5i6j7` is unconditional and concerns two +different indexes; this change does not alter that historical migration. + +Before the new revision is recorded, an explicit migration invocation can opt in: + +```bash +cd apps/api +uv run alembic -x retire_legacy_token_lookup=true upgrade head +``` + +An ordinary deployment records the revision without retirement. Re-running +`upgrade head` with that argument later cannot execute an already-recorded +revision. Use the operator command below to activate at the same head, including +after later migrations. Do not downgrade the application schema for activation. + +## Operator procedure + +Supply `DATABASE_URL` through the protected environment. The command also honors +`DB_SSL_MODE`, `DB_SSL_CERT`, `DB_SSL_KEY`, and `DB_SSL_ROOT_CERT`; absent explicit +values, connection URL and libpq defaults apply. Do not put credentials in shell +arguments or reports. Run commands from the repository root with workspace +dependencies installed. + +1. Confirm the intended database and the application versions needed for rollback. + For production, obtain explicit approval for this index change before applying + it. The investigation and preparation do not authorize a production mutation. +2. Inspect the catalog without changing indexes: + + ```bash + uv run --package knowhere-api-app python scripts/retire_legacy_token_lookup.py + ``` + +3. Explicitly retire the legacy lookup: + + ```bash + uv run --package knowhere-api-app python scripts/retire_legacy_token_lookup.py --apply + ``` + +4. Review the returned catalog. The legacy lookup must be absent; the primary key, + binary lookup, unit index, and unit lookup must remain valid and ready. Check a + real publication and retrieval before deciding to retain the change. Observe + token persistence timing and database I/O; an HTTP success alone does not + establish an improvement. + +The helper verifies the binary and unit-leading lookup definitions and their +`indisvalid` and `indisready` states before removing anything. It also validates +the legacy definition before deletion. A same-name partial, incompatible, or +invalid index fails closed. Repeating apply is safe when the legacy index is +already absent and both replacements remain compatible. Only the named legacy +index is targeted; no table data or uniqueness constraint is removed. + +Run one maintenance operation at a time and coordinate other schema changes. +Independent migration and operator connections use concurrent DDL. The operator +sets a five-second lock timeout and a 30-minute statement timeout. Adjust the +latter with `--statement-timeout-seconds` if the rebuild needs a longer window. +Concurrent DDL can still wait for transactions and consume substantial I/O. +Externally owned migration transactions use regular DDL for compatibility with +contract fixtures; do not wrap production index maintenance in such a transaction. + +## Rollback + +Restore the exact old nonunique index independently of the Alembic head: + +```bash +uv run --package knowhere-api-app python scripts/retire_legacy_token_lookup.py --restore +``` + +This creates concurrently: + +```sql +CREATE INDEX idx_document_map_unit_tokens_lookup +ON public.document_map_unit_tokens USING btree (channel, token_hash, map_unit_id); +``` + +The command reads back and verifies the definition, validity, and readiness. +Repeating restore with a compatible valid index is a no-op. Downgrading the new +revision also restores the index, but is unnecessary for this operator rollback. +Application rollback alone does not restore a dropped index. + +Rebuilding scans the token table and consumes storage, CPU, I/O, and WAL; it may +take substantially longer than dropping the index. Allow space for the index and +temporary build work, and account for replica lag. If a concurrent operation is +interrupted, inspect the catalog. An invalid same-name index is deliberately +rejected, not silently reused or automatically deleted. Resolve that catalog +state as a separately reviewed operator action, then retry restoration. Until +the rebuild succeeds, the previous index configuration has not been restored. + +## Verification boundary + +Disposable PostgreSQL contracts cover default retention, opt-in migration, +independent and externally owned transactions, operator activation after the +revision is already recorded, idempotent apply/restore, exact restoration, and +refusal when replacement indexes are missing, incompatible, invalid, or not +ready. They also reject an incompatible legacy index. No production index is +changed by these tests or by adding this revision. diff --git a/docs/design/publication-runtime-failure-diagnosis.md b/docs/design/publication-runtime-failure-diagnosis.md new file mode 100644 index 000000000..e54bccd69 --- /dev/null +++ b/docs/design/publication-runtime-failure-diagnosis.md @@ -0,0 +1,217 @@ +# Publication runtime failure diagnosis + +Evidence collected on 2026-09-30. Production inspection was read-only. Local +publication used a disposable PostgreSQL 15.17 database, Redis, LocalStack, and +the real API and Worker entry points. + +## Production materialization timeout: confirmed cause + +The failed SpaceX demo request started at `2026-09-30T03:07:58.245481Z` and +raised `TimeoutError` at `03:09:16.328Z`. The production API log stream was +`api/api/88eebaaa69a4479e811bbd3fd0a62d95` in `/ecs/knowhere-api-prod`. + +Performance Insights for Aurora instance `knowhere-database-prod-instance-1` +provides direct attribution: + +- The statement was `COPY document_map_unit_tokens (...) FROM STDIN (FORMAT binary)`. +- It appeared in 30 consecutive one-second samples, from `03:08:47Z` through + `03:09:16Z`. +- Filtering by that SQL identifier, all 30 samples were `IO:DataFileRead`. +- No lock wait or CPU sample was attributed to this statement. Total sampled + database load in the surrounding interval never exceeded one active session. + +The immediate cause was waiting for database pages to be read from storage until +COPY exhausted the API's 30-second command budget. This establishes storage read +pressure, rather than an application deadlock or slow Python serialization, as +the dominant observed wait. Sampling does not identify the exact relation or +index supplying each read, and it does not prove that no shorter other waits +occurred between samples. + +CloudWatch provides supporting context at one-minute resolution: + +| Metric | 03:07 UTC | 03:08 UTC | 03:09 UTC | +| --- | ---: | ---: | ---: | +| Serverless capacity, ACU | 0.5 | 0.742 | 2.0 | +| Buffer cache hit ratio | 100% | 99.88% | 88.93% | +| Read IOPS | 0 | 39.2 | 413.7 | +| Read latency | 0 | 0.885 ms | 0.817 ms | +| CPU utilization | 23.9% | 39.9% | 40.0% | + +These observations support a cold working set and storage reads during scale-up. +They do not isolate scale-up as the sole cause. Current configuration reports +0–4 ACU and a 300-second auto-pause interval; it must not be substituted for +historical configuration. Historical ACU utilization reached 100% at 2 ACU. + +## Current production token storage + +A subsequent read-only catalog query estimated 13,276,691 token rows. The heap +occupied 3,394,764,800 bytes and total relation storage was 15,540,740,096 bytes. +All five indexes were valid and ready: + +| Index | Leading keys | Bytes | +| --- | --- | ---: | +| `document_map_unit_tokens_pkey` | `id` | 1,468,669,952 | +| `idx_document_map_unit_tokens_lookup` | `channel, token_hash, map_unit_id` | 4,612,505,600 | +| `idx_document_map_unit_tokens_token_lookup_binary` | `channel, decode(token_hash)` | 1,591,910,400 | +| `idx_document_map_unit_tokens_unit` | `map_unit_id, channel` | 285,138,944 | +| `idx_document_map_unit_tokens_unit_lookup` | `map_unit_id, channel, token_hash` | 4,184,621,056 | + +The binary and unit-lookup indexes include token/frequency payload. The total +index footprint is approximately 12.14 decimal GB. Index maintenance is a +plausible source of the observed reads; Performance Insights alone does not +identify which index dominated. These sizes are not a measured bloat estimate. +This catalog snapshot was taken after production recovery; it does not establish the index inventory or sizes +at the time of the failed candidate request. + +## Real local API result + +An HTTP request to the actual `main.py` API materialized `demo-spacex-s1` with +`candidate` enabled and the unchanged 30-second database timeouts. It completed +with HTTP 200 in 46.289 seconds. The document was then archived and checked. + +The resulting input and index contained 227 chunks, 227 map units, and 10,778 +token rows. The demo input consists of one 1,467,990-character page chunk, 94 image +chunks, and 132 table chunks. It differs from the older 922-chunk benchmark input. + +A separate read-only observer sampled the database every 250 ms: + +- The publication transaction lasted approximately 8.6 seconds. +- Token COPY appeared in two samples spanning approximately 0.3 seconds and + finished in under one second. Samples included `DataFileRead` and `WALWrite`. +- Chunk COPY lasted approximately 1.4 seconds. +- No blocking PID was observed. +- Much of the request time preceded the publication transaction. + +The production timeout did **not** recur against this almost-empty local +relation. This successful result does not establish performance against the +production token/index footprint. + +## Validation gaps and repair choices + +The old benchmark created standalone SQLAlchemy engines with `NullPool`, without +the API's 30-second `statement_timeout` and asyncpg `command_timeout`. Its sync +owner also bypassed Worker startup and the gevent/psycogreen callback. Those gaps +explain why it could not establish deployment compatibility. + +Before this repair, the candidate wrote up to 100,000 token rows in one COPY +command. Baseline uses 5,000-row INSERT batches. This changes the amount of work governed by one +statement timeout even when both paths share the same enclosing transaction. + +Worker compatibility and bounded COPY are implemented and verified below. +Remaining performance investigation: + +1. Compare that change against current batching on a disposable large clone, + using the actual API, production index definitions, the same SpaceX demo, + 30-second timeouts, and database wait sampling. Use a separate cold clone per + variant rather than running one after another against warmed indexes. +2. Attribute index read cost on that clone before changing index inventory or + row ordering. Existing sorting by `channel, token_hash, map_unit_id` already + favors both token-leading indexes. Sorting by map unit would favor the + unit-leading indexes while sacrificing that locality; it has no demonstrated + net benefit yet. + +## Local evidence files + +The following sanitized artifacts were retained in +`/tmp/knowhere-publication-repro/`: + +- `demo-before.log`: actual HTTP result and archive completion. +- `demo-before-db-summary.json`: local database samples for that request. +- `database-observation.jsonl`: captured local wait/progress samples. +- `prod-failure-pi-waits.json`: surrounding database wait samples. +- `prod-failure-pi-sql.json`: SQL attribution in the same interval. +- `prod-failure-token-copy-waits.json`: token COPY filtered wait samples. +- `prod-failure-cloudwatch.json`: capacity, cache, CPU, and I/O metrics. +- `prod-current-token-indexes.json`: current relation sizes and index definitions. + +Production evidence was collected through read-only inspection. No production +configuration or data was changed. + +## Implemented repair and verification + +Worker startup installs `psycogreen`'s process-wide wait callback. Psycopg2 +[does not support COPY with a registered wait callback](https://www.psycopg.org/docs/advanced.html#support-for-coroutine-libraries). +Both chunk and token persistence now use bounded `execute_values` INSERTs on the +owning connection when that callback is installed. Compatible psycopg2 and +asyncpg connections retain COPY. The repair never disables the global callback, +opens another connection, or commits independently. Worker INSERT fallback does +not retain all the performance characteristics of native COPY. + +Both persistence helpers now bound each statement to 1,000 rows. Token splitting +applies within a single large map unit as well as between map units. Chunk SQL +tracing records each actual batch separately. This limits work covered by each +command timeout while retaining a single publication transaction. Neither the +API's 30-second command timeout nor its 30-second statement timeout was raised. +A sufficiently slow batch can still time out; this is not an absolute latency +or availability guarantee. + +The actual local Worker, launched through `apps/worker/worker.py`, first failed +an HTTP-created Markdown ingestion job with the production error: +`copy_expert cannot be used with an asynchronous callback`. After the callback +compatibility fix, the same HTTP/upload/confirm-upload/Celery flow completed in +4.282 seconds, with a 397.991 ms publication transaction. Its two chunks were +visible through HTTP and the document was then archived. This small ingestion +checks Worker runtime compatibility, not full PDF parser throughput. + +Focused database contracts cover native psycopg2, the actual psycogreen callback, +asyncpg, escaped text/JSON/NULL values, callback preservation, and rollback. A +slow-write contract deliberately delays PostgreSQL token writes proportionally +to inserted rows and retains a short per-command budget. The original batching +failed with `TimeoutError`; bounded batches completed and rolling back the owner +transaction removed all 2,501 rows. This is controlled fault injection, not a +simulation of Aurora's exact storage/cache behavior. + +### Real API timeout fault injection + +A temporary statement trigger in the isolated local database adds 3.5 ms of +wait per inserted token. It uses the statement's transition table to calculate +one server-side delay, without changing application code or database timeouts. +The real SpaceX API request under the original batching returned HTTP 500 after +64.453 seconds and logged `TimeoutError` in `asyncpg.copy_records_to_table`. +With bounded batches and exactly the same injected wait, the request returned +HTTP 200 after 75.295 seconds. Publication took 42.434 seconds, including 38.191 +seconds of token persistence split across 11 commands; the longest token command +was 3.543 seconds. The document contained all 227 chunks and was archived. The +failed request left zero document, section, and chunk rows in its namespace. + +This proves that the change avoids the reproduced per-command timeout without +relaxing the timeout or splitting the transaction. It does not show faster disk +I/O: total artificial delay was intentionally unchanged. The delay trigger and +function were removed from the disposable database after this run. Production +was not modified. + +### Final run without injected delay + +The final code was loaded by restarting both real services. SpaceX demo +materialization returned HTTP 200 in 38.236 seconds, including a 5.646-second +publication transaction and 824 ms token persistence (11 token commands, +maximum SQL duration 91 ms). The 227 chunks were readable, classic retrieval +returned the published document, and archive succeeded. This single small, +previously exercised local database is not a cold production-scale comparison. +The HTTP duration also includes bundle preparation outside the publication +transaction; it remains above ten seconds. + +A concurrent HTTP-created Worker job completed successfully (worker-reported +processing duration 4.932 seconds; publication 247 ms); retrieval and archive +also passed. Its client flow took 37.479 seconds while the API was simultaneously +materializing the demo, so that wall time is not an isolated Worker throughput +measurement. Agentic retrieval and a new full PDF parse were not exercised. + +All seven focused contracts passed. Ruff, targeted Pyright, and `git diff +--check` passed. The API and Worker used isolated local services. No production rollout was +performed. + +Additional reproduction artifacts in `/tmp/knowhere-publication-repro/`: +`repro-http.py`, `slow-token-writes.py`, `copy-red.log`, `timeout-red.log`, +`demo-timeout-red.log`, `demo-timeout-green.log`, `demo-final.log`, +`worker-final-result.log`, and `final-contracts.log`. The temporary slowdown +script rejects non-local database endpoints. Environment and authentication +files in that directory contain local credentials and must not be shared. + +Re-run the focused contracts from `apps/api`: + +```sh +../../.venv/bin/python -m pytest \ + tests/contract/test_publication_chunk_copy_contract.py \ + tests/contract/test_publication_token_copy_contract.py -q --tb=short +``` diff --git a/docs/design/publication-scale-cold-cache-findings.md b/docs/design/publication-scale-cold-cache-findings.md new file mode 100644 index 000000000..95cb02753 --- /dev/null +++ b/docs/design/publication-scale-cold-cache-findings.md @@ -0,0 +1,173 @@ +# Publication at production scale with a cold database cache + +## Scope and environment + +This investigation follows the runtime repair in PR #448. We ran the actual +API application and Celery Worker against a disposable PostgreSQL clone on the +Mac benchmark host. Redis and S3 were isolated local services. Production was +inspected read-only; no production configuration, data, or indexes were changed. + +The clone was copied from a cleanly stopped template mounted read-only. We +applied current migrations and restored the two missing token indexes before +comparison. Historical catalog evidence from 2026-09-30 02:07:35 UTC confirms +that the five-index configuration also existed before the production failure. + +| Property | Clone configuration | +| --- | --- | +| Token rows, catalog estimate | 12,423,317 | +| Token heap | 2,385,043,456 bytes | +| Initial total token relation including indexes | 8,440,152,064 bytes | +| PostgreSQL | 15.19 | +| Container memory / swap limit | 1 GiB / no additional swap | +| CPU limit | One CPU | +| Shared buffers | 128 MiB | +| Work memory / maintenance work memory | 4 MiB / 64 MiB | +| Effective cache size | 512 MiB, planner estimate only | +| Maximum connections | 40 | +| Durability | fsync, full_page_writes, synchronous_commit enabled | +| I/O timing | Enabled | +| API statement / command timeout | 30 seconds / 30 seconds | + +Before each publication sample, we stopped only the disposable database, used +`POSIX_FADV_DONTNEED` on its data files, and restarted it. This clears PostgreSQL +buffers and requests eviction of the clone's Linux file cache. We did not clear +global host caches or restart the sealed source. macOS/Colima virtual-disk caches +remain uncontrolled. The API and Worker ran on the development machine through +an SSH database tunnel, so application timing includes network round trips. + +Each materialization used the same `demo-spacex-s1` input: 227 chunks and 10,778 +token rows. We captured `pg_stat_statements`, per-index read counters, publication +traces, and sampled wait events. Samples used separate namespaces; successfully +published documents were retrieved and archived. The same clone was reused, +with cache eviction between samples; earlier successful samples leave a small +amount of additional archived data. These are targeted observations, not +statistical latency percentiles or identical snapshots. + +## Batch repair comparison + +Both variants used candidate publication and the same five token indexes: +primary key, `lookup`, `token_lookup_binary`, `unit`, and `unit_lookup`. + +| Measurement | Before PR #448 | After PR #448 | +| --- | ---: | ---: | +| Complete HTTP materialization | 50.943 s | 54.149 s | +| Publication transaction | 10.416 s | 11.376 s | +| Token persistence stage | 3.652 s | 4.228 s | +| Token COPY calls | 1 | 11 | +| Longest token command, application measurement | 3.551 s | 0.482 s | +| Token COPY execution, database measurement | 3.552 s | 3.848 s | +| Token COPY buffer reads | 25,639 | 25,673 | +| Token COPY read time | 3.092 s | 3.148 s | +| Token COPY WAL bytes | 159,152,618 | 160,566,253 | + +Both requests succeeded with classic retrieval and archive. Neither reproduced +the production 30-second timeout. The first sample attributed 12 active samples +to token COPY waiting on `DataFileRead`, consistent with the database timing. + +The batching change bounded individual command duration but did not reduce +storage reads. In this pair, token-stage duration increased approximately 16% +and publication duration approximately 9%. Database execution increased by +approximately 0.30 seconds; additional connection round trips contribute to +the larger application-observed difference. A single pair cannot establish a +stable regression percentage or a production throughput estimate. + +This also corrects the interpretation of the earlier small-database run: +publication below ten seconds did not establish that latency with a large, +cold database. Full HTTP time includes preparation outside publication. + +## Attribution and a focused index experiment + +The first publication produced these per-index buffer-read deltas: + +| Token index | Read blocks | +| --- | ---: | +| Primary key | 148 | +| Legacy `lookup` | 13,435 | +| `token_lookup_binary` | 11,015 | +| `unit` | 400 | +| `unit_lookup` | 1,416 | + +These counters cover the publication observation window and are separate from +COPY's statement-level accounting. Both token-leading indexes dominate reads; +legacy `lookup` accounts for approximately half of observed index reads. + +Current classic discovery and `corpus.recall` use the binary token hash +predicate. Oversized outline/node filtering still has a normal-hash reader in +`scoring/map_lighting.py`, constrained by map-unit IDs and channel. The retained +`unit_lookup` covering index supports that reader. + +Ten read-only production EXPLAIN plans, spanning scopes of 1, 32, 227, 1,000, +and 2,000 map units and two token sets, selected `unit_lookup`. On the clone, +we froze equivalent fixtures and compared exact result counts and SHA-256 +digests before and after removing only the nonunique legacy `lookup` index. +All ten cases matched; every plan continued to use `unit_lookup`. These checks +do not exercise every historical application query or a full agentic episode. + +### Cold publication after removing legacy lookup + +The final sample retained the patched application and all other settings, and +removed only `idx_document_map_unit_tokens_lookup` in the disposable clone. + +| Measurement | Patched, five indexes | Patched, legacy lookup removed | +| --- | ---: | ---: | +| Complete HTTP materialization | 54.149 s | 54.177 s | +| Publication transaction | 11.376 s | 10.182 s | +| Token persistence stage | 4.228 s | 2.624 s | +| Token COPY execution, database measurement | 3.848 s | 1.989 s | +| Longest token command, application measurement | 0.482 s | 0.290 s | +| Token COPY buffer reads | 25,673 | 12,110 | +| Token COPY read time | 3.148 s | 1.618 s | +| Token COPY WAL bytes | 160,566,253 | 85,682,583 | + +Publication, chunk reads, classic retrieval, and archive all succeeded. Token +buffer reads fell approximately 53%, WAL approximately 47%, database COPY time +approximately 48%, and application token-stage duration approximately 38%. +Publication improved approximately 10.5% in this comparison, while full HTTP +latency remained essentially unchanged. Preparation and other stages still +matter, and the publication transaction remains slightly above ten seconds. + +This isolates a concrete write-cost reduction. It supports preparing legacy +lookup retirement as a separate index change while retaining binary lookup, +unit lookup, the unit index, and the primary key. It does not justify removing +binary lookup, changing retrieval semantics, or claiming production latency. +Before a production index change, account for application versions supported +by rollback and the time/storage required to rebuild the old index. The local +query checks do not cover every scope size or historical query shape. + +## Worker verification + +An HTTP-created Markdown ingestion job traversed real upload, confirmation, +Redis/Celery dispatch, gevent Worker startup, parsing, and publication against +the full five-index clone. It completed in 7.707 seconds end to end; Worker +processing reported 6.519 seconds. Both chunks were readable, classic retrieval +returned the document, and archive succeeded. This checks Worker compatibility +against the large database; the small input is not a full-PDF throughput test. +Copied job callback URLs were cleared in the disposable clone before Worker +startup to isolate external delivery. + +## Limits and evidence + +Production's observed token relation occupied approximately 15.54 GB after +recovery, versus 8.44 GB in the rebuilt clone. Row counts are comparable, but +index density, physical size, churn, and storage behavior are different. +The clone's first token COPY averaged approximately 0.121 ms per PostgreSQL +buffer read. Historical Aurora read latency was approximately 0.8 ms at one-minute +resolution; these are different measurement layers and cannot be substituted +into a precise production latency prediction. Aurora scale-up is not simulated. + +Local artifacts are under `/tmp/knowhere-publication-scale/`: before/after HTTP +results, database snapshots, trace summaries, `comparison.json`, index alignment +DDL, historical inventory evidence, and map-lighting parity results. Environment, +authentication, and frozen data fixtures are private local artifacts and are not +included in this document. No production deployment or index retirement is +performed by this investigation. + +## Experiment completion + +Only three large-document publication samples were run: old candidate batching, +patched batching, and patched batching without legacy lookup. There was one +small real Worker job and ten narrowly scoped query parity checks. No broad +benchmark suite or full PDF reparse was started. The disposable database is +retained stopped with the four-index experimental configuration for follow-up; +the source/template are unchanged. Runtime processes and the measurement tunnel +were stopped after capture. diff --git a/packages/shared-python/pyproject.toml b/packages/shared-python/pyproject.toml index 1f849f9ba..f94bb9ebf 100644 --- a/packages/shared-python/pyproject.toml +++ b/packages/shared-python/pyproject.toml @@ -64,6 +64,7 @@ dependencies = [ "typing-extensions==4.14.1", "jieba==0.42.1", "rank-bm25>=0.2.2", + "jsonschema>=4.20.0", # Text processing "blingfire>=0.1.8", diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index 930ae06cd..a51193377 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -111,6 +111,20 @@ class StorageConfig(BaseModel): "ZIP bundles remain single-object uploads." ), ) + DEMO_CANONICAL_BUNDLE_ENABLED: bool = Field( + default=False, + description=( + "Reuse immutable demo ZIP and raw artifacts by content version. " + "Disabling only changes new uploads; existing canonical results remain readable." + ), + ) + DEMO_PUBLICATION_PREPARATION_CACHE_ENABLED: bool = Field( + default=False, + description=( + "Reuse pure tokenization and token frequencies for versioned canonical " + "demo publication, in a bounded process-local cache. Requires canonical bundles." + ), + ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md,.html,.htm", description="Supported file extensions", diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index af3b9ea2f..bd873b2c9 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -242,7 +242,6 @@ class DocumentMapUnit(Base): unit_kind: Mapped[str] = mapped_column(String(32), nullable=False) path_token_count: Mapped[int] = mapped_column(Integer, nullable=False) content_token_count: Mapped[int] = mapped_column(Integer, nullable=False) - term_search_text_lower: Mapped[str] = mapped_column(Text, nullable=False) # Asset presence under this unit's section, after root-asset remount # (``KnowhereProvider._remount_root_assets``). Lets type-scoped queries # (e.g. chunk_types=["image"]) narrow map-unit candidates *before* diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py b/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py index 31cf581ed..6c773c1c8 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py @@ -25,12 +25,15 @@ def build_decision_trace(steps: list[AgentStep]) -> list[DecisionTraceStep]: phase = "finish" if step.tool_name == "finish" else ( "stop" if not step.tool_name else "tool_call" ) + observation: dict[str, object] = {"observation_text": observation_text} + if step.ref_status is not None: + observation["ref_status"] = step.ref_status trace_steps.append( DecisionTraceStep( step_index=step.step_index, agent="agent_explore", phase=phase, - observation={"observation_text": observation_text}, + observation=observation, decision={ "action": step.tool_name or "no_tool_call", "args": step.tool_args, diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/config.py b/packages/shared-python/shared/services/retrieval/agent_explore/config.py index ad41283c2..258dfca1c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/config.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/config.py @@ -8,6 +8,8 @@ from __future__ import annotations +from shared.services.retrieval.agent_tools.registry import REF_ADDRESS_ONE_OF, REF_ADDRESS_RULE + AGENT_EXPLORE_MODEL = "deepseek-v4-flash" # Model for AGENT_EXPLORE_HARNESS=cursor_sdk (harness/cursor_harness.py) — @@ -46,11 +48,23 @@ "items": { "type": "object", "properties": { - "document_id": {"type": "string"}, - "section_path": {"type": "string"}, - "chunk_id": {"type": "string"}, + "document_id": { + "type": "string", + "description": "Document that owns the cited section or chunk.", + }, + "section_path": { + "type": "string", + "description": "Cited section. Omit when chunk_id is set.", + }, + "chunk_id": { + "type": "string", + "description": "Cited chunk. Omit when section_path is set.", + }, }, "required": ["document_id"], + "oneOf": REF_ADDRESS_ONE_OF, + "additionalProperties": False, + "description": REF_ADDRESS_RULE, }, }, "notes": { @@ -62,6 +76,7 @@ }, }, "required": ["refs"], + "additionalProperties": False, } FINISH_TOOL_DESCRIPTION = ( @@ -82,18 +97,24 @@ ## 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. If a call's arguments depend on another call's result, do -not issue them in the same turn. In particular: do not call `corpus.grep` -in the same turn as `corpus.recall`, `corpus.read`, -`corpus.list_documents`, `corpus.outline`, `corpus.node_filter`, or -`corpus.assets` unless `pattern` / `patterns` is already known — those -calls produce the term; a grep with no term is an empty call. Wait for -the result, then grep. `corpus.list_documents` enumerates the entire -namespace: use it only if the user explicitly asks to list or inventory -the corpus's documents, never as the starting step for a content question. -When you have enough evidence, call `{FINISH_TOOL_NAME}` +You are exploring this corpus autonomously to answer one query. Tool +names in this prompt are the registered names (`corpus.read`). This +function-calling loop exposes the same tools with `.` replaced by `_` +(`corpus_read`); use that underscore form when calling. Logs keep the +dotted name. + +Use the tools above to navigate; you may call several tools in one turn +when they are independent. If a call's arguments depend on another +call's result, do not issue them in the same turn. In particular: do +not call `corpus.grep` in the same turn as `corpus.recall`, +`corpus.read`, `corpus.list_documents`, `corpus.outline`, +`corpus.node_filter`, or `corpus.assets` unless `pattern` / `patterns` +is already known — those calls produce the term; a grep with no term is +an empty call. Wait for the result, then grep. `corpus.list_documents` +enumerates the entire namespace: use it only if the user explicitly +asks to list or inventory the corpus's documents, never as the starting +step for a content question. 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 diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py index 3c97e2f33..ffceac8f3 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py @@ -40,6 +40,7 @@ async def dispatch_tool_call( namespace: str, document_scope: DocumentScope = DocumentScope(), budget: ToolBudget | None = None, + query: str = "", ) -> ToolResult: """Run one ``REGISTRY`` tool call against a fresh, call-scoped DB session. @@ -47,7 +48,10 @@ async def dispatch_tool_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``. + into the (call-scoped, already-closed) ``ToolContext``. ``query`` is the + end user's original query for this episode (see ``ToolContext.query``'s + docstring) — both harnesses' ``run_episode`` already receive it, so this + is a pass-through, not a new source of truth. """ try: async with db_factory() as db: @@ -58,6 +62,7 @@ async def dispatch_tool_call( db_factory=db_factory, budget=budget or ToolBudget(), document_scope=document_scope, + query=query, ) return await REGISTRY.dispatch(name, tool_ctx, args) except Exception as exc: # noqa: BLE001 - one broken tool must not kill the episode 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 index 8fd6f5d13..ba2a3b4bd 100644 --- 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 @@ -86,8 +86,11 @@ budget_status_line, cursor_execute_content, finish_refs_from_args, + invalid_finish_message, + read_ref_status, select_episode_refs, tool_message_content, + validate_finish_args, wire_safe_tool_name, ) from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult @@ -199,6 +202,7 @@ def _dispatch_sync( namespace=namespace, document_scope=document_scope, budget=tool_budget, + query=query, ), loop, ) @@ -207,7 +211,11 @@ def _dispatch_sync( 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) + content = tool_message_content( + tool_result, + tool_name=tool_name, + 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 @@ -226,6 +234,7 @@ def _dispatch_sync( elapsed_ms=elapsed_ms, tokens_used_delta=0, tokens_used_total=budget.tokens_used, + ref_status=read_ref_status(tool_name, tool_result), ) ) if tool_name in EVIDENCE_TOOL_NAMES and not tool_result.error: @@ -251,6 +260,25 @@ def execute( ) def finish_execute(args: dict[str, Any], _ctx: Any) -> str: + args = dict(args or {}) + validation_error = validate_finish_args(args) + if validation_error is not None: + with budget_lock: + steps.append( + AgentStep( + step_index=len(steps), + tool_name=FINISH_TOOL_NAME, + tool_args=args, + observation_text=validation_error, + error=validation_error, + elapsed_ms=0, + tokens_used_delta=0, + tokens_used_total=budget.tokens_used, + ) + ) + return json.dumps( + {"status": "error", "error": invalid_finish_message(validation_error)} + ) selected = finish_refs_from_args(args) cited = selected if selected is not None else [] notes = str(args.get("notes") or "") @@ -261,7 +289,7 @@ def finish_execute(args: dict[str, Any], _ctx: Any) -> str: AgentStep( step_index=len(steps), tool_name=FINISH_TOOL_NAME, - tool_args=dict(args or {}), + tool_args=args, observation_text=f"refs={len(cited)} notes={notes!r}", error=None, elapsed_ms=0, diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py index bdb82508b..57124ca85 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py @@ -69,13 +69,21 @@ build_wire_tool_name_map, finish_refs_from_args, https_image_parts, + invalid_finish_message, model_accepts_images, + read_ref_status, select_episode_refs, tool_message_content, + validate_finish_args, 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 ( + REGISTRY, + ToolBudget, + ToolResult, + load_corpus_schema_text, +) # A tool-role message is kept in full for the turn it was produced plus this # many additional turns, then collapsed to a placeholder — see module @@ -129,14 +137,23 @@ def _build_openai_tools() -> tuple[list[dict[str, Any]], dict[str, str]]: return tools, name_map -def _safe_json_loads(raw: str | None) -> dict[str, Any]: - if not raw: - return {} +def _parse_tool_arguments(raw: str | None) -> tuple[dict[str, Any], str | None]: + """Parse one tool call's JSON arguments, or return the parse error. + + Only a missing argument string or an explicit ``{}`` is an empty object; + an empty string or malformed JSON is an error. Either way the error + fails just that one call (a corpus.* tool or ``finish``) and is sent + back to the model — the episode continues. + """ + if raw is None: + return {}, None try: parsed = json.loads(raw) - except (TypeError, ValueError): - return {} - return parsed if isinstance(parsed, dict) else {} + except (TypeError, ValueError) as exc: + return {}, f"invalid JSON arguments: {exc}" + if not isinstance(parsed, dict): + return {}, f"arguments must be a JSON object, got {type(parsed).__name__}" + return parsed, None def _collapse_stale_tool_messages( @@ -252,11 +269,19 @@ async def run_episode( finish_call = next( (tc for tc in tool_calls if tc.function.name == FINISH_TOOL_NAME), None ) + finish_error: str | None = None if finish_call is not None: - args = _safe_json_loads(finish_call.function.arguments) - finish_refs = finish_refs_from_args(args) + args, parse_error = _parse_tool_arguments(finish_call.function.arguments) + if parse_error is None: + parse_error = validate_finish_args(args) + if parse_error is not None and forced_reason is None: + finish_error = parse_error + if finish_call is not None and finish_error is None: + finish_refs = finish_refs_from_args(args) if parse_error is None else None cited = finish_refs if finish_refs is not None else [] - result_notes = str(args.get("notes") or "") + result_notes = ( + parse_error if parse_error is not None else str(args.get("notes") or "") + ) stop_reason = f"budget_{forced_reason}" if forced_reason else "finished" steps.append( AgentStep( @@ -264,7 +289,7 @@ async def run_episode( tool_name=FINISH_TOOL_NAME, tool_args=args, observation_text=f"refs={len(cited)} notes={result_notes!r}", - error=None, + error=parse_error, elapsed_ms=turn_elapsed_ms, tokens_used_delta=turn_tokens, tokens_used_total=budget.tokens_used, @@ -314,20 +339,32 @@ async def run_episode( first_tool_tokens_recorded = False for tc in tool_calls: tool_started = time.perf_counter() - args = _safe_json_loads(tc.function.arguments) + args, parse_error = _parse_tool_arguments(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, - document_scope=document_scope, - budget=tool_budget, - ) + if tc is finish_call: + tool_result = ToolResult( + text="", error=invalid_finish_message(str(finish_error)) + ) + elif parse_error is not None: + tool_result = ToolResult(text="", error=parse_error) + else: + tool_result = await dispatch_tool_call( + canonical_name, + args, + db_factory=db_factory, + user_id=user_id, + namespace=namespace, + document_scope=document_scope, + budget=tool_budget, + query=query, + ) tool_elapsed_ms = int((time.perf_counter() - tool_started) * 1000) - content = tool_message_content(tool_result, max_chars=tool_budget.max_chars) + content = tool_message_content( + tool_result, + tool_name=canonical_name, + max_chars=tool_budget.max_chars, + ) messages.append( {"role": "tool", "tool_call_id": tc.id, "content": content} ) @@ -366,6 +403,7 @@ async def run_episode( ), tokens_used_delta=0 if first_tool_tokens_recorded else turn_tokens, tokens_used_total=budget.tokens_used, + ref_status=read_ref_status(canonical_name, tool_result), ) ) first_tool_tokens_recorded = True diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/shared.py b/packages/shared-python/shared/services/retrieval/agent_explore/shared.py index d16d06efd..47fbc60ef 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/shared.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/shared.py @@ -18,7 +18,11 @@ from dataclasses import dataclass from typing import Any +import jsonschema +import jsonschema.validators + from shared.services.retrieval.agent_explore.budget import EpisodeBudget +from shared.services.retrieval.agent_explore.config import FINISH_TOOL_SCHEMA from shared.services.retrieval.agent_tools import ToolResult # Tools whose ToolResult.refs point at evidence the agent has actually looked @@ -30,6 +34,11 @@ {"corpus.read", "corpus.assets", "corpus.query_table"} ) +# The two map-narrowing tools bound their own text at MAP_TOOL_CHAR_BUDGET by +# folding whole subtrees (agent_tools.map_render) or failing the call, so +# tool_message_content never cuts them mid-row. +MAP_TOOL_NAMES = frozenset({"corpus.outline", "corpus.node_filter"}) + def model_accepts_images(model: str) -> bool: """OpenAI-compatible models attach images only when the name contains vision.""" @@ -88,31 +97,25 @@ def build_wire_tool_name_map(names: list[str]) -> dict[str, str]: return {wire_safe_tool_name(name): name for name in names} -def tool_message_content(result: ToolResult, *, max_chars: int) -> str: +def tool_message_content(result: ToolResult, *, tool_name: str, 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. + 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. ``MAP_TOOL_NAMES`` are returned + whole: they already fit ``MAP_TOOL_CHAR_BUDGET`` or failed the call. """ if result.error: return f"error: {result.error}" text = result.text or "(empty result)" - if len(text) <= max_chars: + if tool_name in MAP_TOOL_NAMES or 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]" + "or use a more specific ref and call again if you need the rest]" ) @@ -164,6 +167,51 @@ def select_episode_refs( ) +def validate_finish_args(args: dict[str, Any]) -> str | None: + """Validate ``finish``'s own args against ``FINISH_TOOL_SCHEMA``; ``None`` = valid. + + Both harnesses route ``finish`` argument validation through this one + function instead of relying solely on the provider enforcing the closed + schema at generation time — whether either provider actually does that + is unverified (see ``config.py``'s ``FINISH_TOOL_SCHEMA`` docstring + context), so this is the real enforcement point. + """ + if not isinstance(args, dict): + return f"finish: arguments must be a JSON object, got {type(args).__name__}" + validator_cls = jsonschema.validators.validator_for(FINISH_TOOL_SCHEMA) + validator = validator_cls(FINISH_TOOL_SCHEMA) + errors = sorted( + validator.iter_errors(args), key=lambda error: [str(p) for p in error.path] + ) + if not errors: + return None + error = errors[0] + message = error.message + if error.validator == "oneOf" and isinstance(error.schema, dict): + message = str(error.schema.get("description") or message) + location = "".join(f"[{p}]" if isinstance(p, int) else f".{p}" for p in error.path) + return f"finish: invalid arguments{f' at {location}' if location else ''}: {message}" + + +def read_ref_status(tool_name: str, result: ToolResult) -> list[dict[str, Any]] | None: + """``corpus.read``'s per-ref ok/failed list for ``AgentStep.ref_status``.""" + if tool_name != "corpus.read": + return None + statuses = result.payload.get("refs") + return list(statuses) if isinstance(statuses, list) else None + + +def invalid_finish_message(error: str) -> str: + """Model-facing text for a rejected ``finish`` call; the episode goes on.""" + return ( + f"{error}. finish was not accepted and the episode continues. Correct " + 'form: {"refs": [{"document_id": "...", "section_path": "..."} or ' + '{"document_id": "...", "chunk_id": "..."}], "notes": "..."} — each ' + "ref names exactly one of section_path or chunk_id. Call finish again " + "with fixed arguments, or keep exploring with the corpus tools." + ) + + def finish_refs_from_args(args: dict[str, Any] | None) -> list[dict[str, Any]] | None: """Return cited refs, or ``None`` when finish omitted the ``refs`` key. diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/types.py b/packages/shared-python/shared/services/retrieval/agent_explore/types.py index d12cf6b6b..7d16092d7 100644 --- a/packages/shared-python/shared/services/retrieval/agent_explore/types.py +++ b/packages/shared-python/shared/services/retrieval/agent_explore/types.py @@ -18,6 +18,9 @@ class AgentStep: elapsed_ms: int tokens_used_delta: int tokens_used_total: int + # corpus.read only: one {ref, status, reason?} entry per requested ref, + # kept whole in the trace (observation_text is capped there). + ref_status: list[dict[str, Any]] | None = None @dataclass 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 4f98b0a30..225848378 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 @@ -85,12 +85,13 @@ text, a literal marker: **Format trap**: `` inside the marker is written verbatim by the parser and stored as-is — it is the on-disk path (`"///..."`, plain `/`, filename - included), **not** the DB `section_path` you get back from `outline` / - `node_filter` / `recall` (which is `" / "`-joined and excludes the - filename). Do not string-match the marker directly against a DB - `section_path`. Convert it first — `section_path_from_chunk_path()` in - `search/lexical_text.py` already does this conversion and is the function - to reuse when implementing marker resolution, not a new one. + included), **not** the DB `section_path` you get back from + `corpus.outline` / `corpus.node_filter` / `corpus.recall` (which is + `" / "`-joined and excludes the filename). Do not string-match the + marker directly against a DB `section_path`. Convert it first — + `section_path_from_chunk_path()` in `search/lexical_text.py` already + does this conversion and is the function to reuse when implementing + marker resolution, not a new one. ## 3. Asset chunks: `image` / `table`, and `connect_to` @@ -105,9 +106,12 @@ list on the **body chunk**, not a location on the asset: the owning section). This link is **one-directional** (body → asset). There is no stored -asset → body back-link; to find which section(s) an asset belongs to, use -the reverse lookup on the `assets` tool rather than assuming the asset chunk -itself names its host. +asset → body back-link. `corpus.grep`, `corpus.recall`, and `corpus.assets` +resolve the host through that body link and print the hosting +`section_path` on the asset row. If no host is found, the row keeps +`Root` and is marked as having no host section. Use `corpus.assets` +`host_of` to list every host, rather than reading the asset's stored +path. ## 4. Document graph @@ -125,25 +129,88 @@ 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 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 - -| Tool | Use when | Scope | Notes | -|---|---|---|---| -| `list_documents` | The user explicitly asks to list or inventory the corpus's documents | namespace | Returns every document's keywords/summary/type mix/`parse_track`. Do not use as a question-answering cold start. | -| `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. Each line also prints the full `section_path` — copy that value into `read`, not the title alone. Depth-limited by argument, not by a token budget. | -| `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 | Call args are `document_ids` + a `predicates` array; each predicate is `{field: "path"|"summary", terms: [...], match: "substring"|"regex"}` — `field` picks section-path vs. summary matching (not body text). Predicates AND together across the array; `terms` within one predicate OR together. 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 / identifier / number lookup against published `term_search_text` (body text, image descriptions, table summaries/keywords, plus filename and section path) | scoped by document/section/chunk_type | Call args are `pattern` (one term) and/or `patterns` (a list of terms OR'd together) — pass every candidate term/synonym you want to check in one call's `patterns` list rather than issuing several parallel `corpus.grep` calls in the same turn. Returns match count plus snippets. Table and image hits include rendered content; body hits include any connected table or image. Each hit includes `chunk_type` and `document_id`. Body hits (`text`/`page`) also include `section_path` for `read`. `image`/`table` hits also include `chunk_id` — their `section_path` is the document Root, not the host section; `read` those with `chunk_id`, or use `assets` `host_of` to find the host. Grep does not scan table-cell HTML. `document_ids` only narrows the search scope. Never pass `file_path` as `chunk_id`. | -| `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. Table and image hits include rendered content; body hits include any connected table or image. Each hit includes `chunk_type` and `document_id`. Body hits include `section_path` for `read`. `image`/`table` hits include `chunk_id` — their `section_path` is the document Root, not the host section; `read` those with `chunk_id`, or use `assets` `host_of`. | -| `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` / image `file_path` into URLs. Small tables (`< 50` rows and `< 50` columns, and under the observation character budget) return their full HTML. Large tables do not: `read` says the table is too large and lists every row/column header — use `query_table` to select specific cells. 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. | -| `query_table` | `read` said a table is too large and you need selected cells | one table chunk | One read-only `SELECT` against that chunk's grid. Requires `document_id`, `chunk_id`, and `sql`. The SQL table name is `t`; columns are `row_header` plus the table's column headers. Writes, `ATTACH`, and multiple statements are rejected. A missing `LIMIT` is set to 50. Small tables can be queried too, but `read` should show those in full first. | -| `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. Forward lines include `document_id`, `chunk_id`, `file_path`, and the asset's own `section_path` (Root). Addresses only — open the body with `read` (and `query_table` for a large table). `read` an asset with `chunk_id`, not `file_path`. | -| `neighbors` | You need related documents in the same namespace | one document | Document-level `related` edges only — see §4. | - -General rule: prefer `outline` / `node_filter` to locate before `read`ing -body text; prefer `grep` over `recall` when you know the exact string you -are looking for; only fall back to `recall` for genuinely fuzzy questions. +- **Vector**: not available yet. `recall` today is path + content BM25. + Exact-string lookup is `grep`. + +## 6. Tools and how they work together + +Every tool exists to find the `section_path`s (or asset `chunk_id`s) that +answer the query, then hand them to `corpus.read`. Names below are the +registered names (`corpus.outline`). MCP clients call those names as-is. +Exact call parameters live only in each tool's own schema/description +(ask for that, do not memorize names here) — this section is the +collaboration map: which tools narrow a search space, which produce hits, +and what to do with either kind of result. + +```mermaid +flowchart LR + subgraph MapNarrowing["Narrow a map (overview / structural predicate)"] + outline["corpus.outline"] + node_filter["corpus.node_filter"] + end + subgraph LeafHits["Find hits (exact string / fuzzy / asset listing)"] + grep["corpus.grep"] + recall["corpus.recall"] + assets["corpus.assets"] + end + outline -->|narrows scope for| LeafHits + node_filter -->|narrows scope for| LeafHits + MapNarrowing -->|section_path| read["corpus.read"] + LeafHits -->|section_path or chunk_id| read + read -->|table too large| query_table["corpus.query_table"] + read --> finish + query_table --> finish + list_documents["corpus.list_documents (namespace inventory, standalone)"] + neighbors["corpus.neighbors (related documents, standalone)"] +``` + +**`corpus.outline` and `corpus.node_filter`** are the two map-narrowing +tools: both take one or more scope targets (a whole document, or a section +and everything under it) and return the *same* row shape — every +`section_path` in that scope, indented by level, with title/summary. +`corpus.outline` returns that unconditionally (the outline of what is +there); `corpus.node_filter` returns it only for the branches around a +structural predicate match (marked as a hit), still shown as full context +(ancestors + the entire matched subtree), not a bare list of matches. Use +a returned `section_path` to `corpus.read` a branch, or as the scope of a +`corpus.grep` / `corpus.recall` / `corpus.assets` call. Neither is +callable on a single leaf section with nothing under it; `corpus.read` +that directly instead. + +**`corpus.grep`, `corpus.recall`, and `corpus.assets`** are the leaf-hit +tools: they search *within* a scope (the whole corpus, or one narrowed by +a prior `corpus.outline`/`corpus.node_filter` call) and return rows +pointing at specific sections/chunks, not a map. `corpus.grep` is exact +string (several terms are any-match / OR, rows stay in document order, +not relevance order); `corpus.recall` is fuzzy ranked search; +`corpus.assets` lists image/table chunks by type — on its own it is an +unfiltered listing with no way to judge relevance, so pair it with a +prior `corpus.grep`/`corpus.recall` hit (scope to that hit's section, or +reverse-resolve the hit's own asset references) rather than browsing +every asset in a document. Grep does not scan table-cell HTML — a table +only surfaces from `corpus.grep`/`corpus.recall` via its published +summary/keywords/caption, never by a literal cell value; find the table +via `corpus.assets` (or a hit on the surrounding body text) and use +`corpus.read`/`corpus.query_table` to inspect its cells. Every hit row +names a `document_id` plus either a `section_path` (body hits) or a +`chunk_id` (image/table hits). Image/table rows show the hosting +section, not the asset's stored `Root` path — read whichever address +applies with `corpus.read` next, then decide pick or no-pick for that +ref immediately from `corpus.read`'s own per-ref status. + +**`corpus.read`** is where body content actually gets consumed: it accepts +many refs at once (mixing map-tool section_paths and leaf-hit chunk_ids), +and reports each ref's outcome separately so a failed ref does not hide a +successful one. A table `corpus.read` reports as too large points at +**`corpus.query_table`** for cell-level `SELECT`s over that same table. + +**`corpus.list_documents`** and **`corpus.neighbors`** stand outside this +flow: `corpus.list_documents` is a namespace-wide inventory (only for an +explicit "list/inventory the documents" request, never a cold-start for a +content question); `corpus.neighbors` returns document-level related +edges (§4), unscoped by section. + +General rule: narrow with `corpus.outline`/`corpus.node_filter` before +searching a large corpus; prefer `corpus.grep` over `corpus.recall` when +you know the exact string you are looking for, and fall back to +`corpus.recall` only for genuinely fuzzy questions. diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py b/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py new file mode 100644 index 000000000..8358cdf65 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py @@ -0,0 +1,163 @@ +"""Asset -> hosting body section, shared by grep, recall and assets. + +Image/table chunks are stored under their document's ``Root`` section; the +real placement is ``chunk_metadata.connect_to`` on the body chunk that embeds +them (``CORPUS_SCHEMA.md`` §3). Hosts are keyed by ``(document_id, +asset chunk_id)`` — chunk ids are content hashes and repeat across +documents — and read only from each document's current revision. When an +asset has several hosts they are kept in body order, so every tool that +picks one picks the same one. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +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 +from shared.services.retrieval.agent_tools.scope import ScopeTarget +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") + +AssetKey = tuple[str, str] + + +@dataclass(frozen=True) +class AssetHost: + document_id: str + section_path: str + body_chunk_id: str + source_file_name: str + + +def host_in_scope(scope: list[ScopeTarget], *, document_id: str, section_path: str) -> bool: + """A host is in scope when its section lies in some target's subtree. + + An empty ``scope`` means no narrowing. + """ + if not scope: + return True + for target in scope: + if target.document_id != document_id: + continue + if target.section_path is None: + return True + if section_path == target.section_path or section_path.startswith( + f"{target.section_path} / " + ): + return True + return False + + +async def load_asset_hosts( + ctx: ToolContext, + *, + document_ids: Iterable[str] | None, + asset_ids: Iterable[str] | None = None, +) -> dict[AssetKey, list[AssetHost]]: + """Every current-revision body chunk that connects to an asset. + + ``document_ids=None`` scans every document visible to the caller; + ``asset_ids`` limits the result to those asset chunk ids. + """ + wanted = set(asset_ids) if asset_ids is not None else None + stmt = ( + select( + DocumentChunk.document_id, + DocumentChunk.chunk_id, + DocumentChunk.chunk_metadata, + 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( + Document.user_id == ctx.user_id, + Document.namespace == ctx.namespace, + Document.status == "active", + Document.current_job_result_id == DocumentChunk.job_result_id, + ctx.document_scope.predicate(Document.document_id), + DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES), + ) + .order_by(DocumentChunk.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) + ) + if document_ids is not None: + stmt = stmt.where(Document.document_id.in_(sorted(set(document_ids)))) + + hosts: dict[AssetKey, list[AssetHost]] = {} + for document_id, chunk_id, metadata, section_path, source_file_name in ( + await ctx.db.execute(stmt) + ).all(): + for target_id in iter_connected_target_ids({"chunk_metadata": metadata}): + if wanted is not None and target_id not in wanted: + continue + hosts.setdefault((str(document_id), target_id), []).append( + AssetHost( + document_id=str(document_id), + section_path=str(section_path or ""), + body_chunk_id=str(chunk_id), + source_file_name=str(source_file_name or ""), + ) + ) + return hosts + + +def in_scope_hosts( + hosts: dict[AssetKey, list[AssetHost]], + key: AssetKey, + scope: list[ScopeTarget], +) -> list[AssetHost]: + return [ + host + for host in hosts.get(key, []) + if host_in_scope(scope, document_id=host.document_id, section_path=host.section_path) + ] + + +def hosted_section_path( + hosts: dict[AssetKey, list[AssetHost]], + key: AssetKey, + scope: list[ScopeTarget], + *, + stored_path: Any, +) -> tuple[str, bool]: + """``(section_path, hosted)`` for an asset row: its first in-scope host, + or its own stored (Root) path with ``hosted=False``.""" + candidates = in_scope_hosts(hosts, key, scope) + if candidates: + return candidates[0].section_path, True + return str(stored_path or "Root"), False + + +async def host_paths_for_hits( + ctx: ToolContext, + hits: list[dict[str, Any]], + scope: list[ScopeTarget], +) -> dict[AssetKey, tuple[str, bool]]: + """``(section_path, hosted)`` for every image/table hit, same rule as assets.""" + asset_hits = [ + hit + for hit in hits + if str(hit.get("chunk_type") or "").strip() in ASSET_CHUNK_TYPES + ] + if not asset_hits: + return {} + hosts = await load_asset_hosts( + ctx, + document_ids={str(hit.get("document_id") or "") for hit in asset_hits}, + asset_ids={str(hit.get("chunk_id") or "") for hit in asset_hits}, + ) + resolved: dict[AssetKey, tuple[str, bool]] = {} + for hit in asset_hits: + key = (str(hit.get("document_id") or ""), str(hit.get("chunk_id") or "")) + resolved[key] = hosted_section_path( + hosts, key, scope, stored_path=hit.get("section_path") + ) + return resolved diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/map_render.py b/packages/shared-python/shared/services/retrieval/agent_tools/map_render.py new file mode 100644 index 000000000..fb2b83477 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/map_render.py @@ -0,0 +1,159 @@ +"""Render an ``outline``/``node_filter`` map within ``MAP_TOOL_CHAR_BUDGET``. + +A map that fits is returned whole. An oversized one shrinks in three steps, +each only if the previous did not fit: + +1. Drop every row's ``summary``; the structure (rows, indentation, titles, + paths) stays as is. +2. Score sections against the end user's original query + (``ToolContext.query``) and fold the lowest-scoring subtrees into + placeholder rows (``scoring.map_lighting``). The rows that must stay are + the best-scored chain (outline) or the predicate matches (node_filter), + each with its ancestors. +3. If the rows that must stay still overflow, or scoring is impossible (no + query, no usable map-unit index), the call fails and tells the agent the + scope is too large to map, to run ``corpus.grep``/``corpus.recall``/ + ``corpus.assets`` on that same scope instead. + +The text is never cut mid-row. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from shared.services.retrieval.agent_tools.registry import MAP_TOOL_CHAR_BUDGET, ToolContext +from shared.services.retrieval.agent_tools.snippet import format_row +from shared.services.retrieval.scoring.map_lighting import ( + MapNode, + best_score_ids, + fold_map_nodes, + load_leaf_unit_scores, + pool_scores_to_tree, + render_map_lines, +) + +SUMMARIES_OMITTED_NOTE = " (summaries omitted: map over the size limit)" + + +@dataclass +class MapRender: + text: str = "" + visible_ids: set[str] = field(default_factory=set) + error: str | None = None + + +def _nodes( + order: list[tuple[str, str | None]], + rows_by_id: dict[str, dict[str, Any]], + scores: dict[str, float] | None = None, +) -> list[MapNode]: + return [ + MapNode( + section_id=section_id, + parent_section_id=parent_id, + render=format_row(rows_by_id[section_id]), + depth=rows_by_id[section_id]["depth"], + score=(scores or {}).get(section_id, 0.0), + ) + for section_id, parent_id in order + ] + + +def _too_large(tool: str, chars: int, reason: str) -> MapRender: + return MapRender( + error=( + f"{tool}: the map is {chars} chars even without summaries, over the " + f"{MAP_TOOL_CHAR_BUDGET}-char limit, and {reason}. This scope is too " + "large to show as a map: run corpus.grep, corpus.recall or " + "corpus.assets directly on this same scope instead." + ) + ) + + +async def render_map( + ctx: ToolContext, + *, + tool: str, + order: list[tuple[str, str | None]], + rows_by_id: dict[str, dict[str, Any]], + revision_by_document: dict[str, str], + header: Callable[[int, int], str], + keep_ids: set[str] | None = None, +) -> MapRender: + """``order`` lists ``(section_id, in-map parent id)`` in display order. + + ``header(visible_count, hidden_count)`` renders the first line. + ``keep_ids=None`` means "keep the best-scored chain" (outline); a set + means "keep exactly these rows" (node_filter's matches). Ancestors of + kept rows are always kept. + + When the map is oversized, every row's ``summary`` in ``rows_by_id`` is + blanked in place, so the caller's payload rows match the rendered text. + """ + nodes = _nodes(order, rows_by_id) + text = header(len(nodes), 0) + "\n" + "\n".join(render_map_lines(nodes)) + if len(text) <= MAP_TOOL_CHAR_BUDGET: + return MapRender(text=text, visible_ids={node.section_id for node in nodes}) + + for row in rows_by_id.values(): + row["summary"] = "" + nodes = _nodes(order, rows_by_id) + text = header(len(nodes), 0) + SUMMARIES_OMITTED_NOTE + "\n" + "\n".join( + render_map_lines(nodes) + ) + if len(text) <= MAP_TOOL_CHAR_BUDGET: + return MapRender(text=text, visible_ids={node.section_id for node in nodes}) + + chars = len(text) + if not ctx.query.strip(): + return _too_large(tool, chars, "no user query is available to rank sections") + leaf_scores = await load_leaf_unit_scores( + ctx.db, + revision_by_document=revision_by_document, + section_ids=[section_id for section_id, _ in order], + query=ctx.query, + ) + if leaf_scores is None: + return _too_large( + tool, + chars, + "these documents have no usable relevance index to decide which " + "sections to fold", + ) + + children_by_parent: dict[str, list[str]] = {} + for section_id, parent_id in order: + if parent_id is not None: + children_by_parent.setdefault(parent_id, []).append(section_id) + pooled = pool_scores_to_tree( + children_by_parent=children_by_parent, + roots=[section_id for section_id, parent_id in order if parent_id is None], + leaf_scores=leaf_scores, + ) + if keep_ids is None: + keep_ids = best_score_ids(pooled) + + nodes = _nodes(order, rows_by_id, pooled) + widest_header = header(len(nodes), len(nodes)) + SUMMARIES_OMITTED_NOTE + folded = fold_map_nodes( + nodes, + char_budget=MAP_TOOL_CHAR_BUDGET - len(widest_header) - 1, + keep_ids=keep_ids, + ) + if folded.overflow: + return _too_large( + tool, + chars, + "the sections that must stay do not fit even after folding " + "everything else", + ) + return MapRender( + text=header(len(folded.visible_ids), folded.hidden_count) + + SUMMARIES_OMITTED_NOTE + + "\n" + + "\n".join(folded.lines), + visible_ids=folded.visible_ids, + ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py index 4f9c2c3a7..862531daf 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py @@ -13,11 +13,14 @@ from __future__ import annotations +import logging from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager from dataclasses import dataclass, field from typing import Any +import jsonschema +import jsonschema.validators from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.document_scope import DocumentScope @@ -25,6 +28,36 @@ DbFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] +logger = logging.getLogger(__name__) + +# Char budget for the two map tools (``corpus.outline``, ``corpus.node_filter``) +# once they self-bound via lighting instead of a hard mid-text cut — see +# ``scoring/map_lighting.py`` and each harness's ``tool_message_content`` call. +# Every other tool still uses ``EVIDENCE_TEXT_CHAR_BUDGET`` below. +MAP_TOOL_CHAR_BUDGET = 20_000 + +CHUNK_TYPE_VALUES = ("text", "page", "image", "table") + + +# ``oneOf`` for a ``{document_id, section_path | chunk_id}`` ref item: exactly +# one address. The item schema's ``description`` is the validation message. +REF_ADDRESS_ONE_OF: list[dict[str, Any]] = [ + {"required": ["section_path"], "not": {"required": ["chunk_id"]}}, + {"required": ["chunk_id"], "not": {"required": ["section_path"]}}, +] +REF_ADDRESS_RULE = "each ref needs document_id plus exactly one of section_path or chunk_id" + + +def chunk_types_schema(description: str) -> dict[str, Any]: + """``chunk_types`` argument schema shared by every tool that filters on it.""" + return { + "type": "array", + "items": {"type": "string", "enum": list(CHUNK_TYPE_VALUES)}, + "minItems": 1, + "uniqueItems": True, + "description": description, + } + @dataclass(frozen=True) class ToolBudget: @@ -33,7 +66,7 @@ class ToolBudget: ``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 + default (both ``grep`` and ``recall`` call this argument ``limit``) 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 @@ -53,8 +86,12 @@ class ToolBudget: 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)) + """Clamp a caller-requested row count down to ``budget.max_items``. + + The lower bound is the tool schema's ``minimum`` — a non-positive + request fails validation instead of being raised to 1 here. + """ + return min(requested, budget.max_items) @dataclass @@ -67,6 +104,16 @@ class ToolContext: db_factory: DbFactory budget: ToolBudget = field(default_factory=ToolBudget) document_scope: DocumentScope = DocumentScope() + # The end user's original query for this episode, when the caller is + # ``agent_explore`` (``dispatch.dispatch_tool_call`` threads it through + # from ``run_episode``). Used only by ``corpus.outline``/``corpus.node_filter`` + # to score and fold an oversized map (``scoring/map_lighting.py``) — never + # by hit tools, which already take their own ``query``/``pattern``. + # Empty when a caller (e.g. an external MCP client hitting a single + # ``corpus.*`` tool directly) has no such query: an oversized map that + # dropping summaries does not fit then fails that one call as "too + # large to map" instead of guessing. + query: str = "" @dataclass @@ -123,7 +170,84 @@ async def dispatch( spec = self.get(name) if spec is None: return ToolResult(text="", error=f"unknown tool: {name}") - return await spec.run(ctx, args) + schema_error = validate_tool_args(spec, args) + if schema_error is not None: + return ToolResult(text="", error=schema_error) + try: + return await spec.run(ctx, args) + except Exception as exc: # noqa: BLE001 - one broken call must not end the caller's loop + logger.exception("corpus tool %s raised", name) + return ToolResult( + text="", + error=( + f"{name}: internal error ({type(exc).__name__}: {exc}). Only " + "this call failed — retry with a narrower scope, or use " + "another corpus.* tool." + ), + ) + + +def _describe_error_path(path: Any) -> str: + parts: list[str] = [] + for part in path: + if isinstance(part, int): + parts.append(f"[{part}]") + else: + parts.append(f".{part}" if parts else str(part)) + return "".join(parts) or "" + + +def validate_tool_args(spec: ToolSpec, args: dict[str, Any]) -> str | None: + """Validate ``args`` against ``spec.json_schema``; ``None`` means valid. + + Every tool schema declares ``additionalProperties: false`` at every + object level (registered tools are expected to keep this true — see + ``CORPUS_SCHEMA.md`` §6), so an unknown key anywhere in the argument + tree fails here instead of being silently dropped or ignored downstream. + The returned message names the offending path, the tool's legal + top-level argument names, and — for an unknown-key error — the exact + legal names at that nesting level, so the caller can retry correctly + without re-reading the full schema. + """ + if not isinstance(args, dict): + return f"{spec.name}: arguments must be a JSON object, got {type(args).__name__}" + + validator_cls = jsonschema.validators.validator_for(spec.json_schema) + validator = validator_cls(spec.json_schema) + errors = sorted( + validator.iter_errors(args), key=lambda error: [str(p) for p in error.path] + ) + if not errors: + return None + + error = errors[0] + path = _describe_error_path(error.path) + top_level_args = sorted((spec.json_schema.get("properties") or {}).keys()) + + if error.validator == "additionalProperties": + instance = error.instance if isinstance(error.instance, dict) else {} + schema_here = error.schema if isinstance(error.schema, dict) else {} + legal_here = sorted((schema_here.get("properties") or {}).keys()) + extra = sorted(set(instance) - set(legal_here)) + return ( + f"{spec.name}: unknown argument(s) {extra} at {path} — this " + f"object only accepts: {legal_here}. Retry with only those names." + ) + if error.validator == "required": + return ( + f"{spec.name}: {error.message} at {path}. " + f"Top-level arguments: {top_level_args}." + ) + if error.validator == "oneOf": + rule = error.schema.get("description") if isinstance(error.schema, dict) else None + return ( + f"{spec.name}: invalid argument at {path}: {rule or error.message}. " + f"Top-level arguments: {top_level_args}." + ) + return ( + f"{spec.name}: invalid argument at {path}: {error.message}. " + f"Top-level arguments: {top_level_args}." + ) REGISTRY = ToolRegistry() diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/scope.py b/packages/shared-python/shared/services/retrieval/agent_tools/scope.py new file mode 100644 index 000000000..d0ebc4bf1 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agent_tools/scope.py @@ -0,0 +1,173 @@ +"""Shared ``scope`` parameter: narrow a tool call to one or more documents, +optionally to a section subtree within each. + +Unifies what used to be five different narrowing parameters (``outline``'s +``document_id``/``path_prefix``, and ``document_ids`` on ``node_filter``, +``grep``, ``recall``, ``assets``) into one shape every map/hit tool shares: +``scope: [{document_id, section_path?}]``. Omitting ``section_path`` scopes to +the whole document; a given ``section_path`` scopes to that section and +everything under it. + +A ``section_path`` that does not exist on the document's current revision is +an error — this never silently narrows to "whatever did resolve" or falls +back to the whole document. This is intentionally stricter than +``corpus.read``'s suffix-tolerant ``section_path_lookup.resolve_section_path_anchor``: +that module is a citation convenience for a model that may abbreviate a path +when *reading*; ``scope`` is a caller-declared boundary for *searching*, so +an exact match is required — the caller got this path from a prior +``outline``/``node_filter``/hit row and should copy it verbatim. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy import and_, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import Document, DocumentSection +from shared.services.retrieval.document_scope import DocumentScope +from shared.services.retrieval.search.lexical_text import normalize_section_path + +SCOPE_ITEM_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "document_id": { + "type": "string", + "description": "Document to include in this call's scope.", + }, + "section_path": { + "type": "string", + "description": ( + "Restrict this document to this section and everything " + "under it. Must be the exact section_path from a prior " + "outline/node_filter/grep/recall row — not a partial or " + "ancestor-omitting path. Omit for the whole document." + ), + }, + }, + "required": ["document_id"], + "additionalProperties": False, +} + +SCOPE_SCHEMA: dict[str, Any] = { + "type": "array", + "items": SCOPE_ITEM_SCHEMA, + "minItems": 1, + "description": ( + "One or more {document_id, section_path?} targets to search. " + "Omitting section_path on an item scopes to that whole document; " + "giving one scopes to that section and its full subtree. An " + "unknown document_id or section_path fails the call instead of " + "silently searching less than requested." + ), +} + + +@dataclass(frozen=True) +class ScopeTarget: + document_id: str + job_result_id: str + # ``None`` means the whole document; otherwise the exact, already + # section_path-normalized subtree root. + section_path: str | None + + +async def resolve_scope( + db: AsyncSession, + *, + user_id: str, + namespace: str, + document_scope: DocumentScope, + raw_scope: Any, +) -> tuple[list[ScopeTarget], str | None]: + """Resolve a ``scope`` argument into concrete targets, or an error string. + + Every item must name a document the caller can see and, when given, an + exact ``section_path`` on that document's current revision. Any single + failure fails the whole call. + """ + items = raw_scope if isinstance(raw_scope, list) else [] + if not items: + return [], "scope requires at least one {document_id, section_path?} item" + + document_ids: list[str] = [] + for item in items: + if not isinstance(item, dict): + return [], f"scope item must be an object with document_id, got {item!r}" + document_id = str(item.get("document_id") or "").strip() + if not document_id: + return [], "scope item requires document_id" + document_ids.append(document_id) + + 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") + .where(document_scope.predicate(Document.document_id)) + ) + ) + .scalars() + .all() + ) + revision_by_doc = { + d.document_id: d.current_job_result_id for d in documents if d.current_job_result_id + } + + targets: list[ScopeTarget] = [] + for item in items: + document_id = str(item.get("document_id") or "").strip() + job_result_id = revision_by_doc.get(document_id) + if not job_result_id: + return [], f"unknown document_id: {document_id}" + raw_path = str(item.get("section_path") or "").strip() + if not raw_path: + targets.append(ScopeTarget(document_id, job_result_id, None)) + continue + normalized = normalize_section_path(raw_path) + exists = ( + await db.execute( + select(DocumentSection.section_id) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) + .where(DocumentSection.section_path == normalized) + ) + ).first() + if exists is None: + return [], ( + f"unknown section_path for {document_id}: {normalized!r} — " + "use the exact section_path from a prior outline/node_filter/" + "grep/recall row" + ) + targets.append(ScopeTarget(document_id, job_result_id, normalized)) + return targets, None + + +def scope_document_ids(scope: list[ScopeTarget]) -> list[str]: + return sorted({target.document_id for target in scope}) + + +def scope_orm_clause( + scope: list[ScopeTarget], + *, + document_id_col: Any, + section_path_col: Any, +) -> Any: + """OR of per-target ``(document_id AND subtree)`` clauses for an ORM query.""" + per_target = [] + for target in scope: + clauses: list[Any] = [document_id_col == target.document_id] + if target.section_path is not None: + clauses.append( + or_( + section_path_col == target.section_path, + section_path_col.like(f"{target.section_path} / %"), + ) + ) + per_target.append(and_(*clauses)) + return or_(*per_target) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py b/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py index eeaa68bbf..dd7619f90 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py @@ -1,13 +1,23 @@ -"""Shared hit rendering for ``corpus.grep`` and ``corpus.recall``. +"""Shared hit rendering for every ``corpus.*`` tool. ``build_snippet`` windows body text around the first match: head + first-match window + tail, ``...``-joined, overlap-merged. Only the first match is -windowed. ``format_search_hit_line`` renders the model-visible identifier -line so both tools stay on the same shape. +windowed. + +``build_row`` / ``format_row`` are the one row shape every tool shares — +outline/node_filter (map rows: indented, ``summary``), +grep/recall (hit rows: flat, ``snippet``, ``score``), and assets (asset rows: +``chunk_id`` plus the hosting ``section_path``) all build the same record for +``payload["rows"]`` and render it the same way, so a model reads one row +shape regardless of which tool produced it and copies the same fields +(``document_id`` + ``section_path``/``chunk_id``) into ``corpus.read``. """ from __future__ import annotations +from collections.abc import Mapping +from typing import Any + HIT_CONTEXT_CHARS = 80 HEAD_TAIL_CHARS = 50 @@ -65,31 +75,65 @@ def build_snippet( return "".join(parts) -def format_search_hit_line( +def build_row( *, - source_file_name: object, + kind: str, document_id: object, section_path: object, - snippet: str, - chunk_type: object, + title: object = "", chunk_id: object | None = None, - score: object | None = None, -) -> str: - """Render one grep/recall hit with the fields the model can copy into read. + summary: str = "", + snippet: str = "", + score: float | None = None, + depth: int = 0, + is_hit: bool = False, + hosted: bool | None = None, +) -> dict[str, Any]: + """One row record. Every ``corpus.*`` search/map tool returns these in + ``payload["rows"]`` and renders them with ``format_row``. - Body hits keep ``document_id`` + ``section_path``. Image/table hits also - include ``chunk_id`` because their stored ``section_path`` is the document - Root, not the host section. + ``kind`` is ``section`` (outline/node_filter node), ``text``/``page`` + (body hit), or ``image``/``table`` (asset — its own body chunk or a hit + on one). For ``image``/``table`` rows ``section_path`` is the *hosting* + section, and ``hosted`` says whether a host was found (``False`` keeps + the asset's own ``Root`` path); it is ``None`` for every other kind. + ``summary`` is a map-row field; ``snippet``/``score`` are hit-row + fields; ``depth`` indents a map row under its parent; ``is_hit`` marks a + node_filter predicate match. """ - type_label = str(chunk_type or "").strip() - score_part = f" score={score}" if score is not None else "" - snippet_part = f": {snippet!r}" if snippet else "" - if chunk_id: - return ( - f"- [{type_label}] {source_file_name} ({document_id}) " - f"chunk_id={chunk_id} / {section_path}{score_part}{snippet_part}" - ) - return ( - f"- [{type_label}] {source_file_name} ({document_id}) / " - f"{section_path}{score_part}{snippet_part}" - ) + return { + "kind": kind, + "title": str(title or "").strip(), + "document_id": str(document_id), + "section_path": str(section_path), + "chunk_id": str(chunk_id) if chunk_id else None, + "summary": summary, + "snippet": snippet, + "score": score, + "is_hit": is_hit, + "depth": max(depth, 0), + "hosted": hosted, + } + + +def format_row(row: Mapping[str, Any]) -> str: + """Render one ``build_row`` record as the shared model-visible row.""" + indent = " " * row["depth"] + header = f"{indent}- [{row['kind']}]" + if row["title"]: + header += f" {row['title']}" + header += f" | document_id={row['document_id']} section_path={row['section_path']}" + if row["hosted"] is False: + header += " (no host section)" + if row["chunk_id"]: + header += f" chunk_id={row['chunk_id']}" + if row["score"] is not None: + header += f" score={row['score']}" + if row["is_hit"]: + header += " [Hit]" + lines = [header] + if row["summary"]: + lines.append(f"{indent} summary: {row['summary']}") + if row["snippet"]: + lines.append(f"{indent} snippet: {row['snippet']!r}") + return "\n".join(lines) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py index ea78352c3..90a8621b2 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py @@ -2,14 +2,23 @@ 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. +body section lives in ``chunk_metadata.connect_to`` on the *body* chunk. +Both directions resolve hosts through ``agent_tools.asset_hosts`` (keyed by +document + chunk, current revision only), the same resolver ``corpus.grep`` +and ``corpus.recall`` use: forward rows show the first in-scope host's +section_path, and a ``scope`` section subtree keeps an asset only when one +of its hosts lies in that subtree. An asset with no host keeps ``Root`` and +is marked as unhosted. ``chunk_metadata`` is a plain ``JSON`` column (not ``JSONB``), so a containment query (``@>``) is not available here — that operator is JSONB-only in PostgreSQL. + +Forward search on its own is an unfiltered listing with no way to decide +which assets matter for the current question — pair it with a prior +``corpus.recall``/``corpus.grep`` hit (scope to that hit's section, or use +its ``host_of`` reverse lookup) rather than browsing every asset in a +document. """ from __future__ import annotations @@ -19,32 +28,49 @@ from sqlalchemy import select from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.asset_hosts import ( + hosted_section_path, + in_scope_hosts, + load_asset_hosts, +) 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.agent_tools.scope import ( + SCOPE_SCHEMA, + ScopeTarget, + resolve_scope, + scope_document_ids, +) +from shared.services.retrieval.agent_tools.snippet import build_row, format_row 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." + "section(s) embed or reference them via connect_to. Rows show the " + "hosting section_path; an asset no body section embeds keeps Root " + "and is marked '(no host section)'. A forward search alone only " + "lists candidate assets — it does not tell you which ones matter " + "for the current question. Pair it with a prior corpus.recall/" + "corpus.grep hit: scope this call to that hit's section, or " + "reverse-resolve the hit's own connect_to targets via host_of " + "instead of browsing every asset in a document." ), json_schema={ "type": "object", "properties": { - "document_ids": {"type": "array", "items": {"type": "string"}}, + "scope": SCOPE_SCHEMA, "type": { "type": "string", "enum": ["image", "table", "any"], "default": "any", + "description": "Forward-search asset type. Ignored for host_of.", }, "query": { "type": "string", @@ -52,34 +78,34 @@ }, "host_of": { "type": "array", - "items": {"type": "string"}, + "items": {"type": "string", "minLength": 1}, + "minItems": 1, "description": "Asset chunk_ids to reverse-resolve to hosting sections.", }, }, "required": [], + "additionalProperties": False, }, ) 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() - ] + scope: list[ScopeTarget] = [] + if args.get("scope") is not None: + scope, scope_error = await resolve_scope( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + document_scope=ctx.document_scope, + raw_scope=args.get("scope"), + ) + if scope_error is not None: + return ToolResult(text="", error=f"assets: {scope_error}") host_of = [str(c).strip() for c in (args.get("host_of") or []) if str(c).strip()] - scope_filters: list[Any] = [ - ctx.document_scope.predicate(Document.document_id), - 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 _reverse_lookup(ctx, scope=scope, target_ids=host_of) return await _forward_search( ctx, - scope_filters=scope_filters, + scope=scope, asset_type=str(args.get("type") or "any").strip().lower(), query=str(args.get("query") or "").strip().lower(), ) @@ -88,7 +114,7 @@ async def assets(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: async def _forward_search( ctx: ToolContext, *, - scope_filters: list[Any], + scope: list[ScopeTarget], asset_type: str, query: str, ) -> ToolResult: @@ -98,14 +124,32 @@ async def _forward_search( .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) + .where( + ctx.document_scope.predicate(Document.document_id), + Document.user_id == ctx.user_id, + Document.namespace == ctx.namespace, + Document.status == "active", + Document.current_job_result_id == DocumentChunk.job_result_id, + DocumentChunk.chunk_type.in_(sorted(types)), + ) + .order_by(DocumentChunk.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) ) - rows = (await ctx.db.execute(stmt)).all() + if scope: + stmt = stmt.where(Document.document_id.in_(scope_document_ids(scope))) + asset_rows = (await ctx.db.execute(stmt)).all() - results: list[dict[str, Any]] = [] - for chunk, section_path, source_file_name in rows: + hosts = await load_asset_hosts( + ctx, + document_ids=sorted({chunk.document_id for chunk, _path, _name in asset_rows}), + asset_ids={chunk.chunk_id for chunk, _path, _name in asset_rows}, + ) + section_scoped = any(target.section_path is not None for target in scope) + + rows: list[dict[str, Any]] = [] + for chunk, stored_path, source_file_name in asset_rows: + key = (chunk.document_id, chunk.chunk_id) + if section_scoped and not in_scope_hosts(hosts, key, scope): + continue metadata = chunk.chunk_metadata if isinstance(chunk.chunk_metadata, dict) else {} summary = str(metadata.get("summary") or "").strip() keywords = metadata.get("keywords") or [] @@ -115,84 +159,93 @@ async def _forward_search( ) 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, - } + section_path, hosted = hosted_section_path(hosts, key, scope, stored_path=stored_path) + rows.append( + build_row( + kind=str(chunk.chunk_type or ""), + document_id=chunk.document_id, + section_path=section_path, + title=source_file_name, + chunk_id=chunk.chunk_id, + summary=summary, + hosted=hosted, + ) ) - if len(results) >= ctx.budget.max_items: + if len(rows) >= ctx.budget.max_items: break - lines = [f"assets={len(results)}"] - for r in results: - lines.append( - f"- [{r['chunk_type']}] document_id={r['document_id']} " - f"chunk_id={r['chunk_id']} file_path={r['file_path']} " - f"section_path={r['section_path']} — {r['summary']}" - ) - + lines = [f"assets={len(rows)}", *(format_row(row) for row in rows)] return ToolResult( text="\n".join(lines), - payload={"assets": results}, - refs=[{"document_id": r["document_id"], "chunk_id": r["chunk_id"]} for r in results], + payload={"rows": rows, "details": {}}, + refs=[{"document_id": row["document_id"], "chunk_id": row["chunk_id"]} for row in rows], ) async def _reverse_lookup( ctx: ToolContext, *, - scope_filters: list[Any], + scope: list[ScopeTarget], target_ids: list[str], ) -> ToolResult: - target_set = set(target_ids) - stmt = ( - select(DocumentChunk, DocumentSection.section_path, Document.source_file_name) + hosts = await load_asset_hosts( + ctx, + document_ids=scope_document_ids(scope) if scope else None, + asset_ids=target_ids, + ) + asset_stmt = ( + select(DocumentChunk.document_id, DocumentChunk.chunk_id, DocumentChunk.chunk_type) .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)) + .where( + Document.current_job_result_id == DocumentChunk.job_result_id, + DocumentChunk.document_id.in_(sorted({key[0] for key in hosts})), + DocumentChunk.chunk_id.in_(target_ids), + ) ) - 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, - } + asset_type_by_key = { + (str(document_id), str(chunk_id)): str(chunk_type or "") + for document_id, chunk_id, chunk_type in ( + (await ctx.db.execute(asset_stmt)).all() if hosts else [] + ) + } + + rows: list[dict[str, Any]] = [] + unhosted: list[str] = [] + seen: set[tuple[str, str, str]] = set() + for target_id in target_ids: + keys = sorted(key for key in hosts if key[1] == target_id) + found = False + for key in keys: + for host in in_scope_hosts(hosts, key, scope): + found = True + row_key = (host.document_id, host.section_path, target_id) + if row_key in seen: + continue + seen.add(row_key) + rows.append( + build_row( + kind=asset_type_by_key.get(key, "asset"), + document_id=host.document_id, + section_path=host.section_path, + title=host.source_file_name, + chunk_id=target_id, + hosted=True, + ) ) + if not found: + unhosted.append(target_id) - 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']}" - ) - + lines = [format_row(row) for row in rows] + lines.extend( + f"- {target_id}: no host section found in this scope (asset stays under Root)" + for target_id in unhosted + ) return ToolResult( - text="\n".join(lines) if lines else "no hosts found", - payload={"hosts_by_target": hosts_by_target}, + text="\n".join(lines), + payload={"rows": rows, "details": {"unhosted": unhosted}}, refs=[ - {"document_id": host["document_id"], "section_path": host["section_path"]} - for hosts in hosts_by_target.values() - for host in hosts + {"document_id": row["document_id"], "section_path": row["section_path"]} + for row in rows ], ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py index 961feb81a..87c40fe49 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py @@ -12,9 +12,10 @@ includes rendered table/image content. 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. ++ first-match window + tail, ``...``-joined, overlap-merged). Rows render +through the shared ``agent_tools.snippet.format_row`` — the same row shape +``corpus.outline``/``corpus.node_filter``/``corpus.recall``/``corpus.assets`` +use, so a model reads one shape regardless of which tool produced it. """ from __future__ import annotations @@ -26,22 +27,26 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult +from shared.services.retrieval.agent_tools.asset_hosts import host_paths_for_hits from shared.services.retrieval.agent_tools.explore_mount import mount_explore_hits from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, capped_limit, + chunk_types_schema, register_tool, ) -from shared.services.retrieval.agent_tools.snippet import ( - HIT_CONTEXT_CHARS, - build_snippet, - format_search_hit_line, +from shared.services.retrieval.agent_tools.scope import ( + SCOPE_SCHEMA, + ScopeTarget, + resolve_scope, + scope_orm_clause, ) +from shared.services.retrieval.agent_tools.snippet import build_row, build_snippet, format_row from shared.services.retrieval.settings import ASSET_CHUNK_TYPES -_DEFAULT_MAX_RESULTS = 30 -_DEFAULT_CONTEXT_CHARS = HIT_CONTEXT_CHARS +_DEFAULT_LIMIT = 30 +_DEFAULT_CONTEXT_CHARS = 80 def _terms_from_args(args: dict[str, Any]) -> list[str]: @@ -89,13 +94,17 @@ def _term_search(terms: list[str]) -> tuple[re.Pattern[str], Any]: "Exact string search against published term_search_text " "(body text, image descriptions, table summaries/keywords, plus " "filename and section path). Returns the total number of matching " - "chunks plus a capped list of snippets. Table and image hits include " - "rendered content; body hits include any connected table or image. " - "Provide 'pattern' for one term, or 'patterns' for several candidate " - "terms OR'd together in this single call (e.g. synonyms) — issue one " - "call with multiple terms instead of several parallel corpus.grep " - "calls for different terms in the same turn. At least one of " - "pattern/patterns is required." + "chunks plus a capped list of snippet rows (score-free — this is an " + "exact match, not a ranked search; use corpus.recall for ranking). " + "Table and image hits include rendered content; body hits include " + "any connected table or image — read the hit's document_id + " + "section_path with corpus.read next. Provide 'pattern' for one " + "term, or 'patterns' for several candidate terms OR'd together in " + "this single call (e.g. synonyms) — issue one call with multiple " + "terms instead of several parallel corpus.grep calls for different " + "terms in the same turn. Several terms are any-match (OR), and " + "rows stay in document order — this tool does not rank by " + "relevance. At least one of pattern/patterns is required." ), json_schema={ "type": "object", @@ -107,74 +116,102 @@ def _term_search(terms: list[str]) -> tuple[re.Pattern[str], Any]: "patterns": { "type": "array", "items": {"type": "string"}, - "description": "Several search terms OR'd together in this one call.", + "description": ( + "Several search terms OR'd together in this one call. " + "A chunk matches if any term hits. Results stay in " + "document order, not relevance order." + ), + }, + "scope": SCOPE_SCHEMA, + "chunk_types": chunk_types_schema( + "Restrict hits to these chunk types. Old type names are rejected." + ), + "context_chars": { + "type": "integer", + "minimum": 1, + "default": _DEFAULT_CONTEXT_CHARS, + "description": "Characters of context around the first match in the snippet.", + }, + "limit": { + "type": "integer", + "minimum": 1, + "default": _DEFAULT_LIMIT, + "description": "Max rows to return. The total match count is still reported.", }, - "document_ids": {"type": "array", "items": {"type": "string"}}, - "chunk_types": {"type": "array", "items": {"type": "string"}}, - "context_chars": {"type": "integer", "default": _DEFAULT_CONTEXT_CHARS}, - "max_results": {"type": "integer", "default": _DEFAULT_MAX_RESULTS}, }, "anyOf": [ {"required": ["pattern"]}, {"required": ["patterns"]}, ], + "additionalProperties": False, }, ) async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: terms = _terms_from_args(args) if not terms: return ToolResult(text="", error="grep requires pattern or patterns") - 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() - ] + context_chars = int(args.get("context_chars", _DEFAULT_CONTEXT_CHARS)) + requested_limit = int(args.get("limit", _DEFAULT_LIMIT)) + limit = capped_limit(requested_limit, ctx.budget) chunk_types = { str(t).strip().lower() for t in (args.get("chunk_types") or []) if str(t).strip() } + scope: list[ScopeTarget] = [] + scope_filter = None + if args.get("scope") is not None: + scope, scope_error = await resolve_scope( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + document_scope=ctx.document_scope, + raw_scope=args.get("scope"), + ) + if scope_error is not None: + return ToolResult(text="", error=f"grep: {scope_error}") + scope_filter = scope_orm_clause( + scope, + document_id_col=DocumentChunk.document_id, + section_path_col=DocumentSection.section_path, + ) + compiled, term_filter = _term_search(terms) # Match the indexed haystack first so PostgreSQL can use - # idx_document_chunks_term_trgm. Joining documents first makes the - # planner filter every in-scope chunk and ignore that index. - matched = select( - DocumentChunk.id, - DocumentChunk.chunk_id, - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.chunk_type, - DocumentChunk.term_search_text, - DocumentChunk.content, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentChunk.section_id, - DocumentChunk.sort_order, - ).where( - DocumentChunk.term_search_text.is_not(None), - term_filter, - DocumentChunk.user_id == ctx.user_id, - DocumentChunk.namespace == ctx.namespace, - ctx.document_scope.predicate(DocumentChunk.document_id), + # idx_document_chunks_term_trgm. Section is joined here (not later) so a + # scope's section-subtree condition can apply inside this same CTE. + matched = ( + select( + DocumentChunk.id, + DocumentChunk.chunk_id, + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.chunk_type, + DocumentChunk.term_search_text, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata, + DocumentChunk.section_id, + DocumentChunk.sort_order, + DocumentSection.section_path, + ) + .select_from(DocumentChunk) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .where( + DocumentChunk.term_search_text.is_not(None), + term_filter, + DocumentChunk.user_id == ctx.user_id, + DocumentChunk.namespace == ctx.namespace, + ctx.document_scope.predicate(DocumentChunk.document_id), + ) ) - if document_ids: - matched = matched.where(DocumentChunk.document_id.in_(document_ids)) + if scope_filter is not None: + matched = matched.where(scope_filter) if chunk_types: matched = matched.where( func.lower(DocumentChunk.chunk_type).in_(sorted(chunk_types)) ) matched = matched.cte("matched").prefix_with("MATERIALIZED") - scope_filters = [ - Document.user_id == ctx.user_id, - Document.namespace == ctx.namespace, - Document.status == "active", - Document.current_job_result_id == matched.c.job_result_id, - ctx.document_scope.predicate(Document.document_id), - ] - if document_ids: - scope_filters.append(Document.document_id.in_(document_ids)) - rows_stmt = ( select( matched.c.chunk_id, @@ -186,20 +223,25 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: matched.c.chunk_metadata, matched.c.job_result_id, JobResult.job_id, - DocumentSection.section_path, + JobResult.document_metadata["result_raw_prefix"].as_string(), + matched.c.section_path, Document.source_file_name, func.count().over().label("total_matches"), ) .select_from(matched) .join(Document, Document.document_id == matched.c.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == matched.c.section_id) .outerjoin(JobResult, JobResult.id == Document.current_job_result_id) - .where(*scope_filters) + .where( + Document.user_id == ctx.user_id, + Document.namespace == ctx.namespace, + Document.status == "active", + Document.current_job_result_id == matched.c.job_result_id, + ) .order_by(matched.c.document_id, matched.c.sort_order) - .limit(max_results) + .limit(limit) ) - rows = (await ctx.db.execute(rows_stmt)).all() - total_matches = int(rows[0][-1]) if rows else 0 + matched_rows = (await ctx.db.execute(rows_stmt)).all() + total_matches = int(matched_rows[0][-1]) if matched_rows else 0 results: list[dict[str, Any]] = [] for ( @@ -212,10 +254,11 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: chunk_metadata, job_result_id, job_id, + result_raw_prefix, section_path, source_file_name, _total_matches, - ) in rows: + ) in matched_rows: text = str(term_search_text or "") match = compiled.search(text) snippet = build_snippet( @@ -234,6 +277,7 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "chunk_metadata": chunk_metadata or {}, "job_result_id": job_result_id, "job_id": job_id, + "result_raw_prefix": result_raw_prefix, } ) @@ -242,29 +286,43 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: results, media = await mount_explore_hits( ctx, results, char_budget=ctx.budget.max_chars ) + host_paths = await host_paths_for_hits(ctx, results, scope) + rows: list[dict[str, Any]] = [] lines = [f"total_matches={total_matches} returned={len(results)}"] - if requested_max_results > max_results: + if requested_limit > limit: lines.append(f"note: capped to budget.max_items={ctx.budget.max_items}") for r in results: chunk_type = str(r["chunk_type"] or "").strip() - lines.append( - format_search_hit_line( - source_file_name=r["source_file_name"], - document_id=r["document_id"], - section_path=r["section_path"], - snippet=r["snippet"], - chunk_type=chunk_type, - chunk_id=r["chunk_id"] if chunk_type in ASSET_CHUNK_TYPES else None, - ) + is_asset = chunk_type in ASSET_CHUNK_TYPES + key = (str(r["document_id"]), str(r["chunk_id"])) + section_path, hosted = ( + host_paths[key] if is_asset and key in host_paths else (r["section_path"], None) ) + row = build_row( + kind=chunk_type or "text", + document_id=r["document_id"], + section_path=section_path, + title=r["source_file_name"], + chunk_id=r["chunk_id"] if is_asset else None, + snippet=r["snippet"], + hosted=hosted if is_asset else None, + ) + rows.append(row) + lines.append(format_row(row)) rendered = str(r.get("rendered") or "").strip() if rendered: lines.append(rendered) return ToolResult( text="\n".join(lines), - payload={"total_matches": total_matches, "results": results}, + payload={ + "rows": rows, + "details": { + "total_matches": total_matches, + "capped": requested_limit > limit, + }, + }, 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 index 5c300cf26..330ccb871 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py @@ -32,6 +32,7 @@ "type": "object", "properties": {}, "required": [], + "additionalProperties": False, }, ) async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: 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 index 09607c5ce..e5498ee6e 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py @@ -31,8 +31,14 @@ ), json_schema={ "type": "object", - "properties": {"document_id": {"type": "string"}}, + "properties": { + "document_id": { + "type": "string", + "description": "Document to find related documents for.", + } + }, "required": ["document_id"], + "additionalProperties": False, }, ) async def neighbors(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py index e8739b090..771dd261e 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py @@ -1,11 +1,21 @@ """``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. +``scoring.node_filter_predicates`` (path/summary substring|regex, fields AND +together, terms OR together) — see that module's docstring — but walks +``document_sections`` rows for the requested scope's current revision +instead of the in-memory map-nav tree. + +Unlike a plain filter, the rendered result is not just the matched rows: it +is each matched section's full ancestor chain (context) plus its entire +descendant subtree, in the same map-row shape ``corpus.outline`` uses (see +``agent_tools.snippet``) — this is a map-narrowing tool, the same family as +``outline``, not a search hit list. Matched rows are marked ``[Hit]``. Past +``registry.MAP_TOOL_CHAR_BUDGET`` chars the map shrinks in +``agent_tools.map_render``: summaries are dropped first, then low-scoring +context is folded with every hit kept: a hit and its ancestor chain are +never folded away, but low-relevance branches under a hit can be hidden +behind a placeholder. """ from __future__ import annotations @@ -14,12 +24,20 @@ from sqlalchemy import select -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.agent_tools.map_render import render_map from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, + chunk_types_schema, register_tool, ) +from shared.services.retrieval.agent_tools.scope import ( + SCOPE_SCHEMA, + resolve_scope, + scope_document_ids, +) +from shared.services.retrieval.agent_tools.snippet import build_row from shared.services.retrieval.scoring.node_filter_predicates import ( FieldPredicate, _compile_predicates, @@ -31,59 +49,81 @@ @register_tool( name="corpus.node_filter", description=( - "Deterministic FOR-ALL/EXISTS/ANY/NOT filter over section titles " - "and summaries — not body text (use corpus.grep for that). Call " - "with document_ids plus a predicates array; each predicate is " - "{field: 'path'|'summary', terms: [...], match: 'substring'|'regex'}. " + "Deterministic FOR-ALL/EXISTS/ANY/NOT filter over section titles and " + "summaries — not body text (use corpus.grep for that). Narrows a map " + "the same way corpus.outline does: the result is each matched " + "section's full ancestor chain plus its entire descendant subtree " + "(marked [Hit] where matched), not a bare list of hits. Call with " + "scope plus a predicates array; each predicate is {field: " + "'path'|'summary', terms: [...], match: 'substring'|'regex'}. " "field=path matches the section path; field=summary matches the " "section summary. Predicates AND together across the array; terms " - "within one predicate's 'terms' list OR together. Returns the " - "complete matched set and its count, never a truncated top-K." + "within one predicate's 'terms' list OR together. Use the returned " + "section_paths to corpus.read the ones that matter, or as the scope " + "of a corpus.grep/corpus.recall/corpus.assets call. A result over " + "20,000 chars is shrunk, not cut off: summaries are omitted first; " + "if it is still too long, every [Hit] row and its parents stay and " + "low-relevance context subtrees are hidden behind a 'hidden N " + "nodes' placeholder. If the hits alone do not fit, the call fails: " + "the scope is too large to map, so use corpus.grep, corpus.recall or " + "corpus.assets on the same scope." ), json_schema={ "type": "object", "properties": { - "document_ids": { - "type": "array", - "items": {"type": "string"}, - "minItems": 1, - }, + "scope": SCOPE_SCHEMA, "predicates": { "type": "array", + "minItems": 1, + "description": ( + "AND across this array. Each item is " + "{field, terms, match}: terms OR together." + ), "items": { "type": "object", "properties": { - "field": {"type": "string", "enum": ["path", "summary"]}, - "terms": {"type": "array", "items": {"type": "string"}}, + "field": { + "type": "string", + "enum": ["path", "summary"], + "description": "Match the section path or the section summary.", + }, + "terms": { + "type": "array", + "items": {"type": "string"}, + "description": "Candidate strings; any one may match.", + }, "match": { "type": "string", "enum": ["substring", "regex"], "default": "substring", + "description": "How each term is matched.", }, }, "required": ["field", "terms"], + "additionalProperties": False, }, - "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." - ), }, + "chunk_types": chunk_types_schema( + "Narrow matched sections to those owning a chunk of one of " + "these types (e.g. ['page'] to filter to page-track leaves " + "only). Omit for no narrowing." + ), }, - "required": ["document_ids", "predicates"], + "required": ["scope", "predicates"], + "additionalProperties": False, }, ) 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") + targets, scope_error = await resolve_scope( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + document_scope=ctx.document_scope, + raw_scope=args.get("scope"), + ) + if scope_error is not None: + return ToolResult(text="", error=f"node_filter: {scope_error}") + raw_predicates = args.get("predicates") or [] if not raw_predicates: return ToolResult(text="", error="node_filter requires predicates") @@ -105,57 +145,25 @@ async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: if failed: return ToolResult( text=f"failed_predicates={failed}", - payload={"cardinality": 0, "matched_sections": [], "failed_predicates": failed}, + payload={ + "rows": [], + "details": {"cardinality": 0, "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") - .where(ctx.document_scope.predicate(Document.document_id)) - ) - ) - .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()) - scoped_document_ids = [document_id for document_id, _ in revision_pairs] - scoped_revision_ids = [revision_id for _, revision_id in revision_pairs] - sections = ( - ( - await ctx.db.execute( - select(DocumentSection).where( - DocumentSection.document_id.in_(scoped_document_ids), - DocumentSection.job_result_id.in_(scoped_revision_ids), - ) - ) - ) - .scalars() - .all() - ) - sections = [ - s for s in sections if revision_by_doc.get(s.document_id) == s.job_result_id - ] + document_ids = scope_document_ids(targets) + revision_by_doc = {target.document_id: target.job_result_id for target in targets} chunk_types = { str(t).strip().lower() for t in (args.get("chunk_types") or []) if str(t).strip() } + allowed_section_ids: set[str] | None = None if chunk_types: chunk_rows = await ctx.db.execute( select(DocumentChunk.section_id, DocumentChunk.chunk_type).where( - DocumentChunk.document_id.in_(scoped_document_ids), - DocumentChunk.job_result_id.in_(scoped_revision_ids), + DocumentChunk.document_id.in_(document_ids), + DocumentChunk.job_result_id.in_(list(revision_by_doc.values())), ) ) allowed_section_ids = { @@ -163,44 +171,116 @@ async def node_filter(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: 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]] = [] + rows_by_id: dict[str, dict[str, Any]] = {} + order: list[tuple[str, str | None]] = [] + matched_ids: set[str] = set() 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): + + for target in targets: + section_stmt = ( + select(DocumentSection) + .where(DocumentSection.document_id == target.document_id) + .where(DocumentSection.job_result_id == target.job_result_id) + .order_by(DocumentSection.sort_order, DocumentSection.section_id) + ) + doc_sections = list((await ctx.db.execute(section_stmt)).scalars().all()) + if target.section_path is not None: + doc_sections = [ + s + for s in doc_sections + if s.section_path == target.section_path + or s.section_path.startswith(f"{target.section_path} / ") + ] + by_id = {s.section_id: s for s in doc_sections} + + hit_ids: set[str] = set() + for section in doc_sections: + if allowed_section_ids is not None and section.section_id not in allowed_section_ids: + continue + values = {"path": section.section_path, "summary": section.summary or ""} + if _node_matches(values, compiled): + hit_ids.add(section.section_id) + + if not hit_ids: continue - matched_sections.append( - { - "document_id": section.document_id, - "section_id": section.section_id, - "section_path": section.section_path, - "summary": section.summary or "", - } + matched_ids.update(hit_ids) + matched_doc_ids.append(target.document_id) + + visible_ids: set[str] = set() + for hit_id in hit_ids: + node = by_id.get(hit_id) + while node is not None: + visible_ids.add(node.section_id) + node = by_id.get(node.parent_section_id) if node.parent_section_id else None + hit_path = by_id[hit_id].section_path + for section in doc_sections: + if section.section_path == hit_path or section.section_path.startswith( + f"{hit_path} / " + ): + visible_ids.add(section.section_id) + + base_level = min(s.section_level for s in doc_sections) + + for section in doc_sections: + if section.section_id not in visible_ids: + continue + rows_by_id[section.section_id] = build_row( + kind="section", + document_id=target.document_id, + section_path=section.section_path, + title=section.section_title, + summary=section.summary or "", + depth=section.section_level - base_level, + is_hit=section.section_id in hit_ids, + ) + parent_id = ( + section.parent_section_id + if section.parent_section_id in visible_ids + else None + ) + order.append((section.section_id, parent_id)) + + if not matched_ids: + return ToolResult( + text="hits=0", + payload={"rows": [], "details": {"cardinality": 0, "matched_document_ids": []}}, + refs=[], ) - 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 = [f"{entry['document_id']} / {entry['section_path']}"] - if entry["summary"]: - block.append(f" summary: {entry['summary']}") - lines.append("\n".join(block)) + def header(visible: int, hidden: int) -> str: + line = f"hits={len(matched_ids)}" + if hidden: + line += f" (folded — {hidden} context nodes hidden; every hit is shown)" + return line + + rendered = await render_map( + ctx, + tool="node_filter", + order=order, + rows_by_id=rows_by_id, + revision_by_document=revision_by_doc, + header=header, + keep_ids=matched_ids, + ) + if rendered.error is not None: + return ToolResult(text="", error=rendered.error) + + visible_rows = [ + rows_by_id[section_id] for section_id, _ in order if section_id in rendered.visible_ids + ] return ToolResult( - text="\n".join(lines), + text=rendered.text, payload={ - "cardinality": len(matched_sections), - "matched_sections": matched_sections, - "matched_document_ids": matched_doc_ids, + "rows": visible_rows, + "details": { + "cardinality": len(matched_ids), + "matched_document_ids": matched_doc_ids, + }, }, refs=[ - {"document_id": entry["document_id"], "section_path": entry["section_path"]} - for entry in matched_sections + {"document_id": row["document_id"], "section_path": row["section_path"]} + for row in visible_rows + if row["is_hit"] ], ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py index c20c6d762..9cca417b5 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py @@ -1,160 +1,211 @@ -"""``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. +"""``corpus.outline`` — titles + summaries, no body text. + +Reads ``document_sections`` for one or more scope targets' current +revision, each optionally narrowed to a section subtree. 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. + +A scope target's own row is always included, so calling this on a section +with no descendants (a leaf) would return exactly that one row with no +outline value — rejected up front instead, pointing at ``corpus.read``. + +The rendered map (one row per section, same shape ``corpus.node_filter`` +uses — see ``agent_tools.snippet``) is never truncated mid-text. Past +``registry.MAP_TOOL_CHAR_BUDGET`` chars it shrinks in +``agent_tools.map_render``: summaries are dropped first, then low-scoring +subtrees become placeholders, then the call fails as too large. """ from __future__ import annotations from typing import Any -from sqlalchemy import func, select +from sqlalchemy import select -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import Document, DocumentSection +from shared.services.retrieval.agent_tools.map_render import render_map from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, register_tool, ) -from shared.services.retrieval.search.lexical_text import normalize_section_path +from shared.services.retrieval.agent_tools.scope import ( + SCOPE_SCHEMA, + resolve_scope, +) +from shared.services.retrieval.agent_tools.snippet import build_row @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." + "Map what a document (or one of its sections) covers: every " + "section_path beneath a scope target, any depth, with title + " + "summary, indented by level — no body text. Use this " + "for 'what does this document/section cover' questions, not for " + "finding an answer to a specific fact (use corpus.recall/corpus.grep " + "for that). Accepts several scope targets in one call. Each target " + "must be a non-leaf section (or the whole document, via a scope " + "item with no section_path) — calling this on a section with no " + "children is rejected; read it directly with corpus.read instead. " + "Use the returned section_paths to corpus.read the ones that matter, " + "or as the scope of a corpus.grep/corpus.recall/corpus.assets call. " + "A result over 20,000 chars is shrunk, not cut off: summaries are " + "omitted first; if it is still too long, the lowest-relevance " + "subtrees are hidden behind a 'hidden N nodes' placeholder — " + "re-scope to that section_path to expand it. If even that does not " + "fit, the call fails: the scope is too large to map, so use " + "corpus.grep, corpus.recall or corpus.assets on the same scope." ), 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." - ), - }, + "scope": SCOPE_SCHEMA, "depth": { "type": "integer", + "minimum": 0, "description": ( - "Max levels below path_prefix (or below the document root " - "when path_prefix is omitted) to include. Omit for the " - "full subtree." + "Max levels below each scope target to include. Omit " + "for unlimited depth." ), }, }, - "required": ["document_id"], + "required": ["scope"], + "additionalProperties": False, }, ) 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") + targets, scope_error = await resolve_scope( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + document_scope=ctx.document_scope, + raw_scope=args.get("scope"), + ) + if scope_error is not None: + return ToolResult(text="", error=f"outline: {scope_error}") + 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") - .where(ctx.document_scope.predicate(Document.document_id)) + + document_ids = sorted({target.document_id for target in targets}) + documents = ( + ( + await ctx.db.execute( + select(Document).where(Document.document_id.in_(document_ids)) + ) ) - ).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) + .scalars() + .all() ) - 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 + source_file_name_by_doc = {d.document_id: d.source_file_name or "" for d in documents} + + rows_by_id: dict[str, dict[str, Any]] = {} + order: list[tuple[str, str | None]] = [] + + for target in targets: + section_stmt = ( + select(DocumentSection) + .where(DocumentSection.document_id == target.document_id) + .where(DocumentSection.job_result_id == target.job_result_id) + .order_by(DocumentSection.sort_order, DocumentSection.section_id) ) - if prefix_section is None: - return ToolResult( - text="", error=f"unknown path_prefix for {document_id}: {prefix}" + doc_sections = list((await ctx.db.execute(section_stmt)).scalars().all()) + + if target.section_path is None: + base_level = 0 + scoped = doc_sections + else: + root_section = next( + (s for s in doc_sections if s.section_path == target.section_path), None ) - 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']} " - f"| section_path={node['section_path']} (chunks={node['chunk_count']})" - ) - if node["summary"]: - line += f"\n{indent} summary: {node['summary']}" - lines.append(line) + if root_section is None: + return ToolResult( + text="", + error=( + f"outline: unknown section_path for {target.document_id}: " + f"{target.section_path}" + ), + ) + has_children = any( + s.section_path.startswith(f"{target.section_path} / ") + for s in doc_sections + ) + if not has_children: + return ToolResult( + text="", + error=( + f"outline: {target.section_path} in {target.document_id} is a " + "leaf section (no child sections) — call corpus.read on it " + "directly instead of outline." + ), + ) + base_level = root_section.section_level + scoped = [ + s + for s in doc_sections + if s.section_path == target.section_path + or s.section_path.startswith(f"{target.section_path} / ") + ] + + if depth is not None: + scoped = [s for s in scoped if (s.section_level - base_level) <= depth] + + scoped_ids = {s.section_id for s in scoped} + + for section in scoped: + parent_id = ( + section.parent_section_id + if section.parent_section_id in scoped_ids + else None + ) + rows_by_id[section.section_id] = build_row( + kind="section", + document_id=target.document_id, + section_path=section.section_path, + title=section.section_title, + summary=section.summary or "", + depth=section.section_level - base_level, + ) + order.append((section.section_id, parent_id)) + + if not order: + return ToolResult(text="sections=0", payload={"rows": [], "details": {}}, refs=[]) - text = ( - f"document={document.source_file_name} ({document_id}) " - f"sections={len(nodes)}\n" + "\n".join(lines) + document_names = ", ".join( + f"{source_file_name_by_doc.get(doc_id, doc_id)} ({doc_id})" + for doc_id in document_ids ) + + def header(visible: int, hidden: int) -> str: + line = f"documents={document_names} sections={visible}" + if hidden: + line += f" (folded — {hidden} nodes hidden)" + return line + + rendered = await render_map( + ctx, + tool="outline", + order=order, + rows_by_id=rows_by_id, + revision_by_document={target.document_id: target.job_result_id for target in targets}, + header=header, + ) + if rendered.error is not None: + return ToolResult(text="", error=rendered.error) + + visible_rows = [ + rows_by_id[section_id] for section_id, _ in order if section_id in rendered.visible_ids + ] return ToolResult( - text=text, - payload={"document_id": document_id, "sections": nodes}, + text=rendered.text, + payload={ + "rows": visible_rows, + "details": {"hidden_count": len(order) - len(visible_rows)}, + }, refs=[ - {"document_id": document_id, "section_path": node["section_path"]} - for node in nodes + {"document_id": row["document_id"], "section_path": row["section_path"]} + for row in visible_rows ], ) diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py index 914fecd70..69ce7f63c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py @@ -37,14 +37,21 @@ json_schema={ "type": "object", "properties": { - "document_id": {"type": "string"}, - "chunk_id": {"type": "string"}, + "document_id": { + "type": "string", + "description": "Document owning the table chunk.", + }, + "chunk_id": { + "type": "string", + "description": "The table chunk_id, from a prior read/recall/grep row.", + }, "sql": { "type": "string", "description": "A single SELECT. LIMIT is added when omitted.", }, }, "required": ["document_id", "chunk_id", "sql"], + "additionalProperties": False, }, ) async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: @@ -61,7 +68,10 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: document = ( await ctx.db.execute( - select(Document, JobResult.job_id) + select( + Document, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ) .select_from(Document) .outerjoin(JobResult, JobResult.id == Document.current_job_result_id) .where(Document.document_id == document_id) @@ -73,7 +83,7 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ).first() if document is None: return ToolResult(text="", error=f"unknown document_id: {document_id}") - doc, job_id = document + doc, job_id, result_raw_prefix = document if not doc.current_job_result_id: return ToolResult(text="", error=f"unknown document_id: {document_id}") @@ -100,6 +110,7 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "content": chunk.content, "file_path": chunk.file_path, "job_id": job_id, + "result_raw_prefix": result_raw_prefix, } ) except TableDownloadError as exc: diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py index ea7f946a9..bd3b6632c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py @@ -39,6 +39,8 @@ ) from shared.models.database.job_result import JobResult from shared.services.retrieval.agent_tools.registry import ( + REF_ADDRESS_ONE_OF, + REF_ADDRESS_RULE, ToolContext, ToolResult, register_tool, @@ -63,6 +65,20 @@ _SAME_AS_MARKER_RE = re.compile(r"\[SAME-AS (.+?) p(\d+)\]") _BODY_CHUNK_TYPES = ("text", "page") +# Repeated after the ref list and after the body: a long body can be cut by +# the per-turn text cap, and the ref list alone may scroll out of view. +_PICK_REMINDER = ( + "[pick: decide now, for each [ok] ref, whether to cite it in finish; " + "[failed] refs cannot be cited]" +) + + +def _ref_status_line(entry: dict[str, Any]) -> str: + tag = "[ok]" if entry["status"] == "ok" else f"[failed: {entry['reason']}]" + return ( + f"{tag} document_id={entry['document_id']} " + f"section_path={entry['section_path']} chunk_id={entry['chunk_id']}" + ) def _section_belongs_to_resolved_path( @@ -234,42 +250,6 @@ def _add(row: dict[str, Any]) -> None: return media -def _normalize_read_refs(args: dict[str, Any]) -> list[Any]: - """Accept the canonical ``refs`` list, or a flat single-document shorthand. - - Deterministic, not model-guessing: observed live tool calls sometimes - hoist ``document_id`` to the top level alongside ``section_path(s)`` / - ``chunk_id(s)`` instead of nesting each pair inside ``refs`` — the exact - shape the ``json_schema`` above documents. Rather than relying on the - model to always match the schema, normalize the known equivalent flat - shape here so a well-formed ``document_id`` isn't discarded over an - outer-structure mismatch. Does not change behavior when ``refs`` is - already a non-empty list. - """ - refs = args.get("refs") - if isinstance(refs, list) and refs: - return refs - - document_id = str(args.get("document_id") or "").strip() - if not document_id: - return [] - - normalized: list[dict[str, Any]] = [] - section_path = args.get("section_path") - if isinstance(section_path, str) and section_path.strip(): - normalized.append({"document_id": document_id, "section_path": section_path.strip()}) - for path in args.get("section_paths") or []: - if isinstance(path, str) and path.strip(): - normalized.append({"document_id": document_id, "section_path": path.strip()}) - chunk_id = args.get("chunk_id") - if isinstance(chunk_id, str) and chunk_id.strip(): - normalized.append({"document_id": document_id, "chunk_id": chunk_id.strip()}) - for cid in args.get("chunk_ids") or []: - if isinstance(cid, str) and cid.strip(): - normalized.append({"document_id": document_id, "chunk_id": cid.strip()}) - return normalized - - @register_tool( name="corpus.read", description=( @@ -279,7 +259,9 @@ def _normalize_read_refs(args: dict[str, Any]) -> list[Any]: "large tables as row/column headers plus a pointer to " "corpus.query_table), and converts asset/page_assets references to " "URLs. Use after outline/node_filter/recall/grep have located where " - "to look." + "to look. Each ref's outcome (ok, or failed with a reason) is " + "reported separately — decide pick/no-pick for each ref right " + "after seeing its status; a failed ref defaults to not picked." ), json_schema={ "type": "object", @@ -289,14 +271,26 @@ def _normalize_read_refs(args: dict[str, Any]) -> list[Any]: "items": { "type": "object", "properties": { - "document_id": {"type": "string"}, - "section_path": {"type": "string"}, - "chunk_id": {"type": "string"}, + "document_id": { + "type": "string", + "description": "Document that owns this section or chunk.", + }, + "section_path": { + "type": "string", + "description": "Section to read. Omit when chunk_id is set.", + }, + "chunk_id": { + "type": "string", + "description": "Chunk to read. Omit when section_path is set.", + }, }, "required": ["document_id"], + "oneOf": REF_ADDRESS_ONE_OF, + "additionalProperties": False, + "description": REF_ADDRESS_RULE, }, "minItems": 1, - "description": "Each ref needs document_id and either section_path or chunk_id.", + "description": "Each ref needs document_id and exactly one of section_path or chunk_id.", }, "mode": { "type": "string", @@ -307,16 +301,25 @@ def _normalize_read_refs(args: dict[str, Any]) -> list[Any]: "section_path ref; ignored for chunk_id refs." ), }, - "include_assets": {"type": "boolean", "default": True}, - "resolve_same_as": {"type": "boolean", "default": True}, + "include_assets": { + "type": "boolean", + "default": True, + "description": "Inline connect_to images and tables into the body.", + }, + "resolve_same_as": { + "type": "boolean", + "default": True, + "description": "Replace page-track SAME-AS markers with the owner text.", + }, }, "required": ["refs"], + "additionalProperties": False, }, ) async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: - refs = _normalize_read_refs(args) - if not refs: - return ToolResult(text="", error="read requires refs") + refs = args.get("refs") + if not isinstance(refs, list) or not refs: + return ToolResult(text="", error="read requires a non-empty refs list") mode = str(args.get("mode") or "self").strip().lower() if mode not in ("self", "descendants"): return ToolResult(text="", error=f"unsupported mode: {mode}") @@ -347,19 +350,42 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: 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] = {} + raw_prefix_by_revision: dict[str, str | None] = {} if job_result_ids: job_rows = await ctx.db.execute( - select(JobResult.id, JobResult.job_id).where(JobResult.id.in_(job_result_ids)) + select( + JobResult.id, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ).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} + for revision_id, job_id, raw_prefix in job_rows.all(): + if revision_id and job_id: + job_id_by_revision[str(revision_id)] = str(job_id) + raw_prefix_by_revision[str(revision_id)] = raw_prefix + + # One entry per input ref (by index), in call order — the per-ref + # ok/failed signal this tool now surfaces instead of a single joined + # error string. "pending" entries are finalized once section_job refs' + # matched chunks are known, below. + ref_status: list[dict[str, Any]] = [ + { + "document_id": str(ref.get("document_id") or "").strip(), + "section_path": str(ref.get("section_path") or "").strip() or None, + "chunk_id": str(ref.get("chunk_id") or "").strip() or None, + "status": "pending", + "reason": None, + "chunk_ids": [], + } + for ref in refs + ] emit_items: list[tuple[str, Any]] = [] - errors: list[str] = [] - for ref in refs: + for index, ref in enumerate(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}") + ref_status[index]["status"] = "failed" + ref_status[index]["reason"] = 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() @@ -381,9 +407,12 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ) ).first() if row is None: - errors.append(f"unknown chunk_id: {chunk_id} in {document_id}") + ref_status[index]["status"] = "failed" + ref_status[index]["reason"] = f"unknown chunk_id: {chunk_id} in {document_id}" continue chunk, resolved_section_path = row + ref_status[index]["status"] = "ok" + ref_status[index]["chunk_ids"] = [chunk.chunk_id] emit_items.append( ( "chunk_rows", @@ -392,6 +421,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "document_id": document_id, "job_result_id": job_result_id, "job_id": job_id, + "result_raw_prefix": raw_prefix_by_revision.get(job_result_id), "source_file_name": source_file_name, "chunk_id": chunk.chunk_id, "section_id": chunk.section_id, @@ -407,7 +437,8 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: continue if not section_path: - errors.append(f"ref for {document_id} needs section_path or chunk_id") + ref_status[index]["status"] = "failed" + ref_status[index]["reason"] = f"ref for {document_id} needs section_path or chunk_id" continue resolved_path, path_error = await resolve_section_path_anchor( @@ -417,17 +448,20 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: section_path=section_path, ) if path_error or not resolved_path: - errors.append(path_error or f"unknown section_path for {document_id}") + ref_status[index]["status"] = "failed" + ref_status[index]["reason"] = path_error or f"unknown section_path for {document_id}" continue emit_items.append( ( "section_job", { + "ref_index": index, "document_id": document_id, "job_result_id": job_result_id, "resolved_path": resolved_path, "source_file_name": source_file_name, "job_id": job_id, + "result_raw_prefix": raw_prefix_by_revision.get(job_result_id), }, ) ) @@ -494,6 +528,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: base_rows.extend(payload) continue job = payload + ref_index = job["ref_index"] matched_ids = { section.section_id for section in section_matches @@ -510,14 +545,17 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: for section in section_matches if section.section_id in matched_ids } + job_chunk_ids: list[str] = [] for chunk in batched_chunks: if chunk.section_id not in matched_ids: continue + job_chunk_ids.append(chunk.chunk_id) base_rows.append( { "document_id": job["document_id"], "job_result_id": job["job_result_id"], "job_id": job["job_id"], + "result_raw_prefix": job.get("result_raw_prefix"), "source_file_name": job["source_file_name"], "chunk_id": chunk.chunk_id, "section_id": chunk.section_id, @@ -532,11 +570,28 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "file_path": chunk.file_path, } ) + if job_chunk_ids: + ref_status[ref_index]["status"] = "ok" + ref_status[ref_index]["chunk_ids"] = job_chunk_ids + else: + ref_status[ref_index]["status"] = "failed" + ref_status[ref_index]["reason"] = ( + f"no body chunk found for {job['resolved_path']} in " + f"{job['document_id']} (image/table-only or empty section)" + ) + status_lines = [_ref_status_line(entry) for entry in ref_status] if not base_rows: return ToolResult( - text="", - error="no chunks resolved for given refs" + (f" ({'; '.join(errors)})" if errors else ""), + text="\n".join(status_lines), + payload={"chunks": [], "refs": ref_status}, + error=( + "read: every ref failed, nothing was read:\n" + + "\n".join(status_lines) + + "\nNone of these can be picked for finish. Fix each ref " + "(copy document_id + section_path/chunk_id exactly from an " + "outline/node_filter/grep/recall/assets row) and retry." + ), ) if resolve_same_as_flag: @@ -592,19 +647,18 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: composed["content"] = _image_display_content(row) assembled.append(composed) - lines = [] - if errors: - lines.append(f"errors: {'; '.join(errors)}") + lines = ["refs:", *(f" {line}" for line in status_lines), _PICK_REMINDER] for row in assembled: lines.append( f"### {row.get('source_file_name')} ({row.get('document_id')}) / " f"{row.get('section_path')} [{row.get('chunk_type')}]" ) lines.append(str(row.get("content") or "")) + lines.append(_PICK_REMINDER) return ToolResult( text="\n".join(lines), - payload={"chunks": assembled, "errors": errors}, + payload={"chunks": assembled, "refs": ref_status}, 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 index 7ad3159b2..56ff0c2e5 100644 --- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py +++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py @@ -1,149 +1,43 @@ """``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): +One live scoring path: ``search.map_unit_discovery.map_unit_discovery`` +(persisted map-unit BM25 over path + content, already RRF-fused there). +Exact-string lookup is ``corpus.grep``, not a second channel here. -- ``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. +``scope`` is applied in that same discovery SQL: a named section is that +section and everything under it, scored before ranking. Classic retrieval +does not pass section targets, so its SQL is unchanged. -``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. +Rows render through the shared ``agent_tools.snippet.format_row`` — the +same row shape ``corpus.outline``/``corpus.node_filter``/``corpus.grep``/ +``corpus.assets`` use. """ from __future__ import annotations -from shared.services.retrieval.document_scope import DocumentScope - from typing import Any -from sqlalchemy import select, text -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import DocumentChunk +from shared.services.retrieval.agent_tools.asset_hosts import host_paths_for_hits from shared.services.retrieval.agent_tools.explore_mount import mount_explore_hits from shared.services.retrieval.agent_tools.registry import ( ToolContext, ToolResult, capped_limit, + chunk_types_schema, register_tool, ) -from shared.services.retrieval.agent_tools.snippet import ( - build_snippet, - format_search_hit_line, +from shared.services.retrieval.agent_tools.scope import ( + SCOPE_SCHEMA, + ScopeTarget, + resolve_scope, + scope_document_ids, ) +from shared.services.retrieval.agent_tools.snippet import build_row, build_snippet, format_row from shared.services.retrieval.hydration.row_utils import normalize_chunk_type from shared.services.retrieval.settings import ASSET_CHUNK_TYPES from shared.services.retrieval.search.map_unit_discovery import map_unit_discovery -from shared.services.retrieval.search.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, jr.job_id -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 -JOIN job_results jr ON jr.id = dmu.job_result_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 _term_channel_rows( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - document_scope: DocumentScope, - 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, scope_params = document_scope.sql() - params.update(scope_params) - 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), - "content": chunk.content, - "file_path": chunk.file_path, - "chunk_metadata": chunk.chunk_metadata or {}, - "job_result_id": chunk.job_result_id, - "job_id": unit_row["job_id"], - } - ) - return results +_DEFAULT_LIMIT = 10 def _identifier_snippet(row: dict[str, Any]) -> str: @@ -164,98 +58,88 @@ def _identifier_snippet(row: dict[str, Any]) -> str: 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 chunk_type, " - "document_id, path and snippet. Table and image hits include rendered " - "content; body hits include any connected table or image. Image/table " - "hits also include chunk_id. Call corpus.read using document_id plus " - "section_path or chunk_id, not the filename." + "where the answer lives. Scores path and content with BM25. " + "Returns candidates with chunk_type, document_id, path and snippet. " + "Table and image hits include rendered content; body hits include " + "any connected table or image. Exact string / identifier lookup is " + "corpus.grep, not this tool. Read the candidate's document_id + " + "section_path (or chunk_id, for image/table hits) with corpus.read " + "next, not the filename." ), 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.", + "query": { + "type": "string", + "description": "The fuzzy question to rank path and content against.", + }, + "scope": SCOPE_SCHEMA, + "chunk_types": chunk_types_schema( + "Restrict hits to these chunk types. Old type names are rejected." + ), + "limit": { + "type": "integer", + "minimum": 1, + "default": _DEFAULT_LIMIT, + "description": "Max ranked candidates to return.", }, - "top_k": {"type": "integer", "default": _DEFAULT_TOP_K}, }, "required": ["query"], + "additionalProperties": False, }, ) 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() - ] + requested_limit = int(args.get("limit", _DEFAULT_LIMIT)) + limit = capped_limit(requested_limit, ctx.budget) 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)", - ) - - document_scope = ctx.document_scope.narrow(document_ids) if document_ids else ctx.document_scope - channel_rows: list[list[dict[str, Any]]] = [] - weights: list[float] = [] - if "path_content" in active_channels: - discovery = await map_unit_discovery( + scope: list[ScopeTarget] = [] + if args.get("scope") is not None: + scope, scope_error = await resolve_scope( ctx.db, user_id=ctx.user_id, namespace=ctx.namespace, - query=query, - top_k=top_k, - exclude_document_ids=[], - document_scope=document_scope, - exclude_sections=[], - chunk_types=chunk_types, + document_scope=ctx.document_scope, + raw_scope=args.get("scope"), ) - channel_rows.append(list(discovery.payload.get("fused_rows") or [])) - weights.append(1.0) + if scope_error is not None: + return ToolResult(text="", error=f"recall: {scope_error}") - if "term" in active_channels: - term_rows = await _term_channel_rows( - ctx.db, - user_id=ctx.user_id, - namespace=ctx.namespace, - query=query, - document_scope=document_scope, - chunk_types=chunk_types, - top_k=top_k, - ) - channel_rows.append(term_rows) - weights.append(1.0) + document_scope = ( + ctx.document_scope.narrow(scope_document_ids(scope)) if scope else ctx.document_scope + ) + discovery = await map_unit_discovery( + ctx.db, + user_id=ctx.user_id, + namespace=ctx.namespace, + query=query, + top_k=limit, + exclude_document_ids=[], + document_scope=document_scope, + exclude_sections=[], + chunk_types=chunk_types, + section_targets=[ + (target.document_id, target.section_path) for target in scope + ] + or None, + ) + rows = list(discovery.payload.get("fused_rows") or []) - fused = merge_channels_rrf(channel_rows, weights, top_k) media: list[dict[str, str]] = [] - if fused: - fused, media = await mount_explore_hits( - ctx, fused, char_budget=ctx.budget.max_chars + if rows: + rows, media = await mount_explore_hits( + ctx, rows, char_budget=ctx.budget.max_chars ) + host_paths = await host_paths_for_hits(ctx, rows, scope) - lines = [f"candidates={len(fused)}"] - if len(fused) < 2: + payload_rows: list[dict[str, Any]] = [] + lines = [f"candidates={len(rows)}"] + if len(rows) < 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 @@ -267,36 +151,44 @@ async def recall(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: "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: + if requested_limit > limit: lines.append(f"note: capped to budget.max_items={ctx.budget.max_items}") - for row in fused: - snippet = _identifier_snippet(row) - chunk_type = str(row.get("chunk_type") or "").strip() - lines.append( - format_search_hit_line( - source_file_name=row.get("source_file_name"), - document_id=row.get("document_id"), - section_path=row.get("section_path"), - snippet=snippet, - chunk_type=chunk_type, - chunk_id=( - row.get("chunk_id") if chunk_type in ASSET_CHUNK_TYPES else None - ), - score=row.get("score"), - ) + for hit in rows: + snippet = _identifier_snippet(hit) + chunk_type = str(hit.get("chunk_type") or "").strip() + is_asset = chunk_type in ASSET_CHUNK_TYPES + key = (str(hit.get("document_id") or ""), str(hit.get("chunk_id") or "")) + section_path, hosted = ( + host_paths[key] + if is_asset and key in host_paths + else (hit.get("section_path"), None) + ) + raw_score = hit.get("score") + row = build_row( + kind=chunk_type or "text", + document_id=hit.get("document_id"), + section_path=section_path, + title=hit.get("source_file_name"), + chunk_id=hit.get("chunk_id") if is_asset else None, + snippet=snippet, + score=float(raw_score) if raw_score is not None else None, + hosted=hosted if is_asset else None, ) - rendered = str(row.get("rendered") or "").strip() + payload_rows.append(row) + lines.append(format_row(row)) + rendered = str(hit.get("rendered") or "").strip() if rendered: lines.append(rendered) return ToolResult( text="\n".join(lines), - payload={"candidates": fused, "reserved_channels": sorted(reserved_requested)}, + payload={ + "rows": payload_rows, + "details": {"capped": requested_limit > limit}, + }, refs=[ - {"document_id": row.get("document_id"), "chunk_id": row.get("chunk_id")} - for row in fused + {"document_id": hit.get("document_id"), "chunk_id": hit.get("chunk_id")} + for hit in rows ], media=media, ) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 5f0078b51..5fe28067f 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -9,9 +9,8 @@ from shared.services.retrieval.execution.plan import ( run_retrieval_query as execute_retrieval_query, ) -from shared.services.retrieval.search.scoring import merge_channels_rrf -__all__ = ["merge_channels_rrf", "run_retrieval_query"] +__all__ = ["run_retrieval_query"] async def run_retrieval_query( @@ -25,11 +24,9 @@ async def run_retrieval_query( exclude_sections: list[dict[str, str]], include_document_ids: list[str] | None = None, chunk_types: set[str] | None = None, - signal_paths: list[str] | None = None, - filter_mode: str = "delete", - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, - rerank: bool = False, + signal_paths: list[str] | None = None, + filter_mode: str = "delete", + rerank: bool = False, threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, @@ -49,8 +46,6 @@ async def run_retrieval_query( chunk_types=chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, rerank=rerank, threshold=threshold, internal_recall_k=internal_recall_k, diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index 5b44d5918..ad766f5af 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -71,8 +71,6 @@ def _cache_shape_digest( chunk_types: list[str] | set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, rerank: bool = False, threshold: float = 0.0, internal_recall_k: int | None = None, @@ -91,8 +89,6 @@ def _cache_shape_digest( chunk_types_str, ",".join(sorted(signal_paths or [])), filter_mode, - ",".join(sorted(channels or [])), - str(sorted((channel_weights or {}).items())), str(rerank), str(threshold), str(internal_recall_k), diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index 0baad9a46..2b654a36f 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -43,8 +43,6 @@ async def run_retrieval_query( chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, rerank: bool = False, threshold: float = 0.0, internal_recall_k: int | None = None, @@ -67,8 +65,6 @@ async def run_retrieval_query( chunk_types=chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, rerank=rerank, threshold=threshold, internal_recall_k=internal_recall_k, @@ -123,6 +119,7 @@ async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, An "namespace": request.namespace, "query": request.query, "router_used": "empty_query_filtered", + "failure_reason": "empty query — retrieval was not run", "evidence": [], "evidence_text": "", "answer_text": "", 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 64795c294..f2f2b2e01 100644 --- a/packages/shared-python/shared/services/retrieval/execution/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -28,8 +28,6 @@ class RetrievalQuery: chunk_types: set[str] | None = None signal_paths: list[str] | None = None filter_mode: str = "delete" - channels: list[str] | None = None - channel_weights: dict[str, float] | None = None rerank: bool = False threshold: float = 0.0 internal_recall_k: int | None = None @@ -54,8 +52,6 @@ def from_parameters( chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, filter_mode: str = "delete", - channels: list[str] | None = None, - channel_weights: dict[str, float] | None = None, rerank: bool = False, threshold: float = 0.0, internal_recall_k: int | None = None, @@ -76,8 +72,6 @@ def from_parameters( chunk_types=chunk_types, signal_paths=signal_paths, filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, rerank=rerank, threshold=threshold, internal_recall_k=internal_recall_k, @@ -102,8 +96,6 @@ def build_cache_extra(self) -> dict[str, Any]: "chunk_types": sorted(self.chunk_types) if self.chunk_types else None, "signal_paths": self.signal_paths, "filter_mode": self.filter_mode, - "channels": self.channels, - "channel_weights": self.channel_weights, "rerank": self.rerank, "threshold": self.threshold, "internal_recall_k": self.internal_recall_k, @@ -141,8 +133,6 @@ def build_route_context(self) -> RetrievalRouteContext: chunk_types=self.chunk_types, signal_paths=self.signal_paths, filter_mode=self.filter_mode, - channels=self.channels, - channel_weights=self.channel_weights, rerank=self.rerank, threshold=self.threshold, internal_recall_k=self.internal_recall_k, diff --git a/packages/shared-python/shared/services/retrieval/execution/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py index d04b898b8..cf9b6e648 100644 --- a/packages/shared-python/shared/services/retrieval/execution/route_types.py +++ b/packages/shared-python/shared/services/retrieval/execution/route_types.py @@ -23,8 +23,6 @@ class RetrievalRouteContext: chunk_types: set[str] | None signal_paths: list[str] | None filter_mode: str - channels: list[str] | None - channel_weights: dict[str, float] | None rerank: bool threshold: float internal_recall_k: int | None diff --git a/packages/shared-python/shared/services/retrieval/graph/query_service.py b/packages/shared-python/shared/services/retrieval/graph/query_service.py index 819d99469..f5d542b98 100644 --- a/packages/shared-python/shared/services/retrieval/graph/query_service.py +++ b/packages/shared-python/shared/services/retrieval/graph/query_service.py @@ -199,6 +199,10 @@ async def collect_candidate_chunks( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_result.job_id if job_result else None, + 'result_raw_prefix': ( + (job_result.document_metadata or {}).get('result_raw_prefix') + if job_result else None + ), }) if len(rows) >= top_k: break diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py index 8b679e5ea..bec7ea168 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -7,10 +7,11 @@ from shared.services.retrieval.hydration.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type from shared.services.storage.page_pdf_crop import crop_source_pdf_pages +from shared.services.storage.raw_prefix_arguments import RawPrefixArguments from shared.services.storage.result_storage import get_result_storage AssetUrlValue = str -PagePdfRequestKey = tuple[str, tuple[int, ...]] +PagePdfRequestKey = tuple[str, tuple[int, ...], str] def _normalize_artifact_ref(asset_ref: object) -> str | None: @@ -107,7 +108,7 @@ def _resolve_page_pdf_request(row: dict[str, Any]) -> PagePdfRequestKey | None: normalized_pages = tuple(sorted({page for page in pages if page > 0})) if not normalized_pages: return None - return job_id, normalized_pages + return job_id, normalized_pages, str(row.get("result_raw_prefix") or "") async def _generate_retrieval_asset_url( @@ -122,6 +123,11 @@ async def _generate_retrieval_asset_url( job_id, artifact_ref = request try: return get_result_storage().generate_artifact_url( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, artifact_ref=artifact_ref, ) @@ -146,6 +152,11 @@ async def _generate_page_citation_asset_url( job_id, artifact_ref = request try: return get_result_storage().generate_artifact_url( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, artifact_ref=artifact_ref, ) @@ -176,6 +187,11 @@ async def _enrich_page_assets( if artifact_ref is not None: try: asset_url = get_result_storage().generate_artifact_url( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, artifact_ref=artifact_ref, ) @@ -208,10 +224,15 @@ async def _generate_page_pdf_asset_url_for_request( *, log_context: str, ) -> str | None: - job_id, pages = request + job_id, pages, raw_prefix = request try: return await asyncio.to_thread( crop_source_pdf_pages, + **( + RawPrefixArguments(raw_prefix=raw_prefix) + if raw_prefix + else RawPrefixArguments() + ), job_id=job_id, pages=list(pages), ) diff --git a/packages/shared-python/shared/services/retrieval/hydration/connected.py b/packages/shared-python/shared/services/retrieval/hydration/connected.py index b960946a8..93dc41305 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/connected.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.py @@ -67,11 +67,14 @@ async def hydrate_connected_target_rows( return [] stmt = ( - # Select only the job identifier needed for the public projection. + # Select only job and artifact location fields needed for hydration. # Selecting the JobResult entity triggers its ``chunks`` selectin # relationship, loading the entire legacy job-chunk collection for # every connected revision during final hydration. - select(Document, DocumentChunk, DocumentSection, JobResult.job_id) + select( + Document, DocumentChunk, DocumentSection, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ) .join( DocumentChunk, ( @@ -100,7 +103,7 @@ async def hydrate_connected_target_rows( result = await db.execute(stmt) hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_id in result.all(): + for document, chunk, section, job_id, result_raw_prefix in result.all(): section_path = section.section_path if section else None hydrated_rows.append( { @@ -116,6 +119,7 @@ async def hydrate_connected_target_rows( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_id, + 'result_raw_prefix': result_raw_prefix, 'sort_order': chunk.sort_order, } ) diff --git a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py index c5cb97463..0c1e4430d 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py +++ b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py @@ -26,6 +26,7 @@ TableDownloadError, load_table_html, ) +from shared.services.storage.raw_prefix_arguments import RawPrefixArguments from shared.services.storage.result_storage import get_result_storage _IMAGE_MEDIA_TYPES = { @@ -284,6 +285,11 @@ def _try_read_image_artifact( return None try: temp_path = storage.download_raw_to_temp( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, relative_path=normalized, suffix=Path(normalized).suffix or ".bin", diff --git a/packages/shared-python/shared/services/retrieval/hydration/path.py b/packages/shared-python/shared/services/retrieval/hydration/path.py index b33d193bd..5572f864b 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/path.py +++ b/packages/shared-python/shared/services/retrieval/hydration/path.py @@ -228,6 +228,10 @@ async def _hydrate_chunk_paths( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_result.job_id if job_result else None, + 'result_raw_prefix': ( + (job_result.document_metadata or {}).get('result_raw_prefix') + if job_result else None + ), 'source_chunk_path': chunk.source_chunk_path, 'sort_order': chunk.sort_order, 'hydrate_mode': path_mode, diff --git a/packages/shared-python/shared/services/retrieval/hydration/reference.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py index 753fc66ff..5a8999e40 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/reference.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.py @@ -68,11 +68,14 @@ async def hydrate_referenced_chunk_rows( ) stmt = ( - # Select only the job identifier needed for the public projection. + # Select only job and artifact location fields needed for hydration. # Selecting the JobResult entity would trigger its ``chunks`` selectin # relationship, loading the entire legacy job-chunk collection for # every referenced revision during final hydration. - select(Document, DocumentChunk, DocumentSection, JobResult.job_id) + select( + Document, DocumentChunk, DocumentSection, JobResult.job_id, + JobResult.document_metadata["result_raw_prefix"].as_string(), + ) .join(DocumentChunk, chunk_join) .outerjoin( DocumentSection, DocumentSection.section_id == DocumentChunk.section_id @@ -94,7 +97,7 @@ async def hydrate_referenced_chunk_rows( rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {} rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {} - for document, chunk, section, job_id in result.all(): + for document, chunk, section, job_id, result_raw_prefix in result.all(): row = { "document_id": document.document_id, "chunk_id": chunk.chunk_id, @@ -115,6 +118,7 @@ async def hydrate_referenced_chunk_rows( "chunk_metadata": chunk.chunk_metadata or {}, "job_result_id": chunk.job_result_id, "job_id": job_id, + "result_raw_prefix": result_raw_prefix, "source_chunk_path": chunk.source_chunk_path, "sort_order": chunk.sort_order, } diff --git a/packages/shared-python/shared/services/retrieval/hydration/table_grid.py b/packages/shared-python/shared/services/retrieval/hydration/table_grid.py index 28c6159b4..668d041ea 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/table_grid.py +++ b/packages/shared-python/shared/services/retrieval/hydration/table_grid.py @@ -21,6 +21,7 @@ LARGE_TABLE_AXIS, QUERY_TABLE_NAME, ) +from shared.services.storage.raw_prefix_arguments import RawPrefixArguments from shared.services.storage.result_storage import get_result_storage _TABLE_HTML_RE = re.compile(r")", re.IGNORECASE) @@ -58,6 +59,11 @@ def load_table_html(row: Mapping[str, Any]) -> str: return "" try: temp_path = storage.download_raw_to_temp( + **( + RawPrefixArguments(raw_prefix=str(row["result_raw_prefix"])) + if row.get("result_raw_prefix") + else RawPrefixArguments() + ), job_id=job_id, relative_path=normalized, suffix=".html", diff --git a/packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py b/packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py new file mode 100644 index 000000000..1857dd84a --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/maintenance/legacy_token_lookup.py @@ -0,0 +1,101 @@ +"""Guarded retirement shared by the migration and the operator command.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.engine import Connection + +_LEGACY_INDEX: str = "idx_document_map_unit_tokens_lookup" +_DEFINITIONS: dict[str, str] = { + _LEGACY_INDEX: "(channel, token_hash, map_unit_id)", + "idx_document_map_unit_tokens_token_lookup_binary": ( + "(channel, decode((token_hash)::text, 'hex'::text)) " + "INCLUDE (map_unit_id, token, frequency)" + ), + "idx_document_map_unit_tokens_unit_lookup": ( + "(map_unit_id, channel, token_hash) INCLUDE (token, frequency)" + ), +} + + +class LegacyTokenLookupMaintenance: + """Change only the legacy nonunique index after catalog validation. + + Concurrent operations require an AUTOCOMMIT connection. Transaction-owned + migration contracts may explicitly use regular DDL. Definitions are checked + conservatively so a same-name invalid, partial, or incompatible index cannot + silently satisfy a prerequisite. Keep these historical definitions stable. + """ + + def __init__(self, connection: Connection, *, is_concurrent: bool = True) -> None: + self._connection: Connection = connection + self._is_concurrent: bool = is_concurrent + + def inspect_indexes(self) -> dict[str, dict[str, str | bool]]: + rows = self._connection.execute( + text( + """ + SELECT relation.relname AS name, index.indisvalid AS is_valid, + index.indisready AS is_ready, + pg_get_indexdef(index.indexrelid) AS definition + FROM pg_index AS index + JOIN pg_class AS relation ON relation.oid = index.indexrelid + WHERE index.indrelid = to_regclass('public.document_map_unit_tokens') + ORDER BY relation.relname + """ + ) + ).mappings().all() + return { + str(row["name"]): { + "is_valid": bool(row["is_valid"]), + "is_ready": bool(row["is_ready"]), + "definition": str(row["definition"]), + } + for row in rows + } + + def _require_compatible_index( + self, indexes: dict[str, dict[str, str | bool]], name: str + ) -> None: + state: dict[str, str | bool] | None = indexes.get(name) + expected: str = ( + f"CREATE INDEX {name} ON public.document_map_unit_tokens " + f"USING btree {_DEFINITIONS[name]}" + ) + if ( + state is None + or not state["is_valid"] + or not state["is_ready"] + or state["definition"] != expected + ): + raise RuntimeError(f"Required token index is missing or incompatible: {name}") + + def apply_retirement(self) -> None: + indexes: dict[str, dict[str, str | bool]] = self.inspect_indexes() + for name in _DEFINITIONS: + if name != _LEGACY_INDEX: + self._require_compatible_index(indexes, name) + if _LEGACY_INDEX not in indexes: + return + self._require_compatible_index(indexes, _LEGACY_INDEX) + concurrent_clause: str = "CONCURRENTLY " if self._is_concurrent else "" + self._connection.execute( + text(f"DROP INDEX {concurrent_clause}IF EXISTS public.{_LEGACY_INDEX}") + ) + if _LEGACY_INDEX in self.inspect_indexes(): + raise RuntimeError("Legacy token lookup index remained after retirement") + + def restore_index(self) -> None: + indexes: dict[str, dict[str, str | bool]] = self.inspect_indexes() + if _LEGACY_INDEX in indexes: + self._require_compatible_index(indexes, _LEGACY_INDEX) + return + concurrent_clause: str = "CONCURRENTLY " if self._is_concurrent else "" + self._connection.execute( + text( + f"CREATE INDEX {concurrent_clause}{_LEGACY_INDEX} " + "ON public.document_map_unit_tokens USING btree " + f"{_DEFINITIONS[_LEGACY_INDEX]}" + ) + ) + self._require_compatible_index(self.inspect_indexes(), _LEGACY_INDEX) 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 1e3617658..6674c601b 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -31,6 +31,7 @@ ) from shared.services.retrieval.scoring.score_units import build_score_units from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache from shared.services.retrieval.publication_strategy import resolve_publication_strategy from shared.services.retrieval.publication_token_copy import ( TokenCopyRow, @@ -130,16 +131,20 @@ def replace_document_map_units( if not unit_id or not section_id: continue map_unit_id = f"dmu_{uuid4().hex}" - path_tokens = str(unit.get("path_search_text") or "").split() - content_tokens = str(unit.get("content_search_text") or "").split() - if path_tokens: + path_frequencies = PublicationPreparationCache.count_tokens( + str(unit.get("path_search_text") or "") + ) + content_frequencies = PublicationPreparationCache.count_tokens( + str(unit.get("content_search_text") or "") + ) + path_token_count = sum(path_frequencies.values()) + content_token_count = sum(content_frequencies.values()) + if path_token_count: path_document_count += 1 - path_total_length += len(path_tokens) - if content_tokens: + path_total_length += path_token_count + if content_token_count: content_document_count += 1 - content_total_length += len(content_tokens) - path_frequencies = Counter(path_tokens) - content_frequencies = Counter(content_tokens) + content_total_length += content_token_count path_unit_df.update(path_frequencies.keys()) content_unit_df.update(content_frequencies.keys()) # ``provider.self_units`` already reflects root-asset remount (assets @@ -160,11 +165,8 @@ def replace_document_map_units( "unit_id": unit_id, "section_id": section_id, "unit_kind": str(unit.get("kind") or "leaf"), - "path_token_count": len(path_tokens), - "content_token_count": len(content_tokens), - "term_search_text_lower": str( - unit.get("term_search_text") or "" - ).lower(), + "path_token_count": path_token_count, + "content_token_count": content_token_count, "has_image": "image" in section_types, "has_table": "table" in section_types, "sort_order": sort_order, diff --git a/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py b/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py index a54f6c23a..d86218b02 100644 --- a/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py +++ b/packages/shared-python/shared/services/retrieval/publication_chunk_copy.py @@ -8,11 +8,15 @@ from datetime import datetime, timezone from typing import Any, Protocol, cast +from psycopg2.extensions import cursor as PsycopgCursor +from psycopg2.extensions import get_wait_callback +from psycopg2.extras import execute_values from sqlalchemy.orm import Session from shared.models.database.document import DocumentChunk from shared.services.jobs.lifecycle.publication_trace_sql import record_publication_sql +_COPY_BATCH_SIZE: int = 1_000 _COLUMNS: tuple[str, ...] = ( "id", "chunk_id", "user_id", "namespace", "document_id", "job_result_id", "section_id", "chunk_type", "content", "content_lexical_text", "path_lexical_text", @@ -100,18 +104,27 @@ def _copy_psycopg( connection: _PsycopgConnection, records: list[tuple[_CopyValue, ...]], ) -> None: - buffer = _encode_csv_rows(records) cursor = connection.cursor() try: - cursor.copy_expert(_COPY_SQL, buffer) + if get_wait_callback() is not None: + # psycogreen uses a process-wide callback that COPY cannot support. + # Keep cooperative I/O and the caller's transaction intact. + execute_values( + cast(PsycopgCursor, cursor), + "INSERT INTO document_chunks (" + ", ".join(_COLUMNS) + ") VALUES %s", + records, + page_size=_COPY_BATCH_SIZE, + ) + else: + buffer = _encode_csv_rows(records) + buffer.seek(0) + cursor.copy_expert(_COPY_SQL, buffer) finally: cursor.close() -def insert_chunks_with_copy(db: Session, chunks: list[DocumentChunk]) -> None: - """COPY prepared chunk objects inside the caller's transaction.""" - if not chunks: - return +def _insert_chunk_batch(db: Session, chunks: list[DocumentChunk]) -> None: + """Persist one bounded statement and record its actual SQL duration.""" connection = db.connection() raw_connection = connection.connection driver = connection.dialect.driver @@ -161,3 +174,9 @@ def insert_chunks_with_copy(db: Session, chunks: list[DocumentChunk]) -> None: is_manual_write=True, write_row_count=len(chunks) if did_succeed else None, ) + + +def insert_chunks_with_copy(db: Session, chunks: list[DocumentChunk]) -> None: + """Persist prepared chunks in bounded batches inside the caller's transaction.""" + for start in range(0, len(chunks), _COPY_BATCH_SIZE): + _insert_chunk_batch(db, chunks[start : start + _COPY_BATCH_SIZE]) diff --git a/packages/shared-python/shared/services/retrieval/publication_preparation_cache.py b/packages/shared-python/shared/services/retrieval/publication_preparation_cache.py new file mode 100644 index 000000000..529c7ef63 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_preparation_cache.py @@ -0,0 +1,143 @@ +"""Bounded reuse of pure lexical preparation during canonical demo publication.""" + +from __future__ import annotations + +import sys +import threading +from collections import Counter, OrderedDict +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass +from hashlib import sha256 + +from shared.utils.text_utils import tokenize_for_retrieval + +_PREPARATION_FORMAT_VERSION = 1 +_MAX_BYTES = 64 * 1024 * 1024 +_MAX_ENTRIES = 4096 +_Scope = tuple[str, str, int, int] +_Key = tuple[_Scope, str, bool, bytes] +_active_scope: ContextVar[_Scope | None] = ContextVar( + "demo_publication_preparation_scope", default=None +) + + +@dataclass(frozen=True) +class _PreparedValue: + tokens: tuple[str, ...] = () + frequencies: tuple[tuple[str, int], ...] = () + size_bytes: int = 0 + + +class PublicationPreparationCache: + """Reuse token sequences and frequencies without storing owner-specific rows. + + Only the canonical materializer opens a scope. Ordinary ingestion, backfills + and retrieval call the same authoritative tokenizer without cache access. + The exact text digest and tokenization mode are part of every key; source + and algorithm versions isolate changed artifacts and preparation semantics. + """ + + _entries: OrderedDict[_Key, _PreparedValue] = OrderedDict() + _size_bytes: int = 0 + _lock: threading.RLock = threading.RLock() + + @classmethod + @contextmanager + def scope( + cls, *, source_id: str, content_version: str, index_format_version: int + ) -> Iterator[None]: + if not source_id or not content_version: + raise ValueError("Demo preparation reuse requires source and content versions") + scope: _Scope = ( + source_id, content_version, index_format_version, _PREPARATION_FORMAT_VERSION + ) + token: Token[_Scope | None] = _active_scope.set(scope) + try: + yield + finally: + _active_scope.reset(token) + + @classmethod + def tokenize(cls, text: str, *, dedupe: bool = False) -> list[str]: + key: _Key | None = cls._build_key(text, operation="tokens", dedupe=dedupe) + cached: _PreparedValue | None = cls._get(key) + if cached is not None: + return list(cached.tokens) + tokens: list[str] = tokenize_for_retrieval( + text, stopwords=[], dedupe=dedupe, min_token_length=2 + ) + if key is not None: + immutable_tokens: tuple[str, ...] = tuple(tokens) + size_bytes: int = sys.getsizeof(immutable_tokens) + sum( + sys.getsizeof(item) for item in immutable_tokens + ) + cls._store(key, _PreparedValue(tokens=immutable_tokens, size_bytes=size_bytes)) + return tokens + + @classmethod + def count_tokens(cls, text: str) -> Counter[str]: + key: _Key | None = cls._build_key(text, operation="frequencies", dedupe=False) + cached: _PreparedValue | None = cls._get(key) + if cached is not None: + return Counter(dict(cached.frequencies)) + frequencies: Counter[str] = Counter(text.split()) + if key is not None: + immutable_frequencies: tuple[tuple[str, int], ...] = tuple(frequencies.items()) + size_bytes: int = sys.getsizeof(immutable_frequencies) + sum( + sys.getsizeof(item) + sys.getsizeof(item[0]) + sys.getsizeof(item[1]) + for item in immutable_frequencies + ) + cls._store( + key, + _PreparedValue(frequencies=immutable_frequencies, size_bytes=size_bytes), + ) + return frequencies + + @classmethod + def clear(cls) -> None: + with cls._lock: + cls._entries.clear() + cls._size_bytes = 0 + + @staticmethod + def _build_key(text: str, *, operation: str, dedupe: bool) -> _Key | None: + scope: _Scope | None = _active_scope.get() + if scope is None: + return None + return scope, operation, dedupe, sha256(text.encode("utf-8")).digest() + + @classmethod + def _get(cls, key: _Key | None) -> _PreparedValue | None: + if key is None: + return None + with cls._lock: + value: _PreparedValue | None = cls._entries.get(key) + if value is not None: + cls._entries.move_to_end(key) + return value + + @classmethod + def _store(cls, key: _Key, value: _PreparedValue) -> None: + # Include conservative per-entry overhead for keys, scope and LRU nodes. + size_bytes: int = value.size_bytes + 1024 + sum( + sys.getsizeof(part) for part in key[0] + ) + if size_bytes > _MAX_BYTES: + return + value = _PreparedValue( + tokens=value.tokens, frequencies=value.frequencies, size_bytes=size_bytes + ) + with cls._lock: + previous: _PreparedValue | None = cls._entries.pop(key, None) + if previous is not None: + cls._size_bytes -= previous.size_bytes + while cls._entries and ( + cls._size_bytes + size_bytes > _MAX_BYTES + or len(cls._entries) >= _MAX_ENTRIES + ): + _, evicted = cls._entries.popitem(last=False) + cls._size_bytes -= evicted.size_bytes + cls._entries[key] = value + cls._size_bytes += size_bytes diff --git a/packages/shared-python/shared/services/retrieval/publication_token_copy.py b/packages/shared-python/shared/services/retrieval/publication_token_copy.py index bfad464e2..83aadf5fd 100644 --- a/packages/shared-python/shared/services/retrieval/publication_token_copy.py +++ b/packages/shared-python/shared/services/retrieval/publication_token_copy.py @@ -8,13 +8,16 @@ from collections.abc import Iterator, Sequence from typing import Protocol, cast +from psycopg2.extensions import cursor as PsycopgCursor +from psycopg2.extensions import get_wait_callback +from psycopg2.extras import execute_values from sqlalchemy.orm import Session from shared.services.jobs.lifecycle.publication_trace_sql import record_publication_sql -# Keep the production-shaped publication in one COPY while retaining a bound -# for larger documents so token persistence does not add avoidable round trips. -_COPY_BATCH_SIZE = 100_000 +# Each COPY has the API's per-command timeout. Bound the work even for a +# single token-heavy map unit; all batches remain in the caller's transaction. +_COPY_BATCH_SIZE: int = 1_000 _TOKEN_COLUMNS = ( "id", "map_unit_id", "channel", "token", "token_hash", "frequency", ) @@ -64,11 +67,21 @@ def _copy_psycopg( connection: _PsycopgConnection, records: list[tuple[str | int, ...]], ) -> None: - buffer = _encode_binary_rows(records) - buffer.seek(0) cursor = connection.cursor() try: - cursor.copy_expert(_COPY_SQL, buffer) + if get_wait_callback() is not None: + # psycogreen uses a process-wide callback that COPY cannot support. + # Keep cooperative I/O and the caller's transaction intact. + execute_values( + cast(PsycopgCursor, cursor), + "INSERT INTO document_map_unit_tokens (" + ", ".join(_TOKEN_COLUMNS) + ") VALUES %s", + records, + page_size=_COPY_BATCH_SIZE, + ) + else: + buffer = _encode_binary_rows(records) + buffer.seek(0) + cursor.copy_expert(_COPY_SQL, buffer) finally: cursor.close() diff --git a/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py index ca73845ba..7e864e831 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py +++ b/packages/shared-python/shared/services/retrieval/scoring/knowhere_hybrid.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Dict, List, Mapping, Optional, Sequence, Tuple -from shared.utils.text_utils import tokenize_for_retrieval as _tokenize_word_level +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache RRF_K = 60 CHANNEL_WEIGHT_PATH = 1.0 @@ -30,12 +30,7 @@ def tokenize_for_retrieval(text: str, *, dedupe: bool = True) -> List[str]: ``min_token_length=2`` so map-unit indexes match ``document_chunks`` search text. Character-level regex tokenization was removed. """ - return _tokenize_word_level( - text, - stopwords=[], - dedupe=dedupe, - min_token_length=2, - ) + return PublicationPreparationCache.tokenize(text, dedupe=dedupe) def tokenize_query_for_ranker(query: str) -> List[str]: diff --git a/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py b/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py new file mode 100644 index 000000000..ced48f737 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/scoring/map_lighting.py @@ -0,0 +1,383 @@ +"""Score-based folding for an oversized ``corpus.outline``/``corpus.node_filter`` map. + +Adapted from the archived MAP-NAV pooling/hiding algorithm +(``deprecated/mapnav/nav/nav_map_scores.py``'s ``_pool_unit_scores_to_tree`` +and ``nav_projection.py``'s ``_apply_budget_hide``) onto the live, +``document_sections``-based tree ``outline``/``node_filter`` already build — +not onto the retired ``ToolSpace``/``LazyKnowhereProvider`` abstraction those +files used. The deprecated files are unmodified; this is a fresh module using +the same idea against live data. + +Only invoked when a rendered map exceeds +``agent_tools.registry.MAP_TOOL_CHAR_BUDGET``. Scoring needs the end user's +original query (``ToolContext.query``) — see that field's docstring for why a +caller without one gets a "narrow the scope" error instead of a silently +unscored fold. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from hashlib import sha256 + +from sqlalchemy import select, tuple_ +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.models.database.document import ( + DocumentMapUnit, + DocumentMapUnitIndex, + DocumentMapUnitToken, +) +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.scoring.persisted_score_load import ( + build_channel_bm25_stats, + combine_average_idf, +) + +_SCORE_CHANNELS = ("path", "content") + + +async def load_leaf_unit_scores( + db: AsyncSession, + *, + revision_by_document: dict[str, str], + section_ids: list[str], + query: str, +) -> dict[str, float] | None: + """BM25 path+content score, pooled onto ``section_id``, for one ``query``. + + Scoped to exactly the ``section_ids`` the caller already rendered (an + outline/node_filter result), not a document/signal-path predicate — the + map tools already know their own scope, so this skips rebuilding it in + SQL. Mirrors ``search.map_unit_discovery``'s scoring formula (same stats + builder, same channel-sum via ``score_persisted_corpus_many``) but reads + the persisted map-unit index directly through the ORM instead of the + discovery module's hand-rolled ``text()`` CTEs, since lighting's scope is + already a concrete id list, not a filter to compile. + + Returns ``None`` when the map cannot be scored at all: the query has no + rankable token, no rendered section has a map unit on its current + revision, or some revision lacks a compatible map-unit index; callers fail the call + instead of folding unscored. A section that simply has no unit of its + own is absent from the mapping and scores ``0.0``. + """ + query_tokens = tokenize_query_for_ranker(query) + if not query_tokens or not section_ids or not revision_by_document: + return None + token_hashes = [sha256(token.encode("utf-8")).hexdigest() for token in query_tokens] + + unit_rows = ( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.section_id.in_(section_ids) + ) + ) + ) + .scalars() + .all() + ) + unit_rows = [ + row + for row in unit_rows + if revision_by_document.get(row.document_id) == row.job_result_id + ] + if not unit_rows: + return None + + map_unit_ids = [row.id for row in unit_rows] + token_rows = await db.execute( + select( + DocumentMapUnitToken.map_unit_id, + DocumentMapUnitToken.channel, + DocumentMapUnitToken.token, + DocumentMapUnitToken.frequency, + ).where( + DocumentMapUnitToken.map_unit_id.in_(map_unit_ids), + DocumentMapUnitToken.channel.in_(_SCORE_CHANNELS), + DocumentMapUnitToken.token_hash.in_(token_hashes), + ) + ) + frequencies: dict[tuple[str, str], dict[str, int]] = {} + for map_unit_id, channel, token, frequency in token_rows.all(): + frequencies.setdefault((str(map_unit_id), str(channel)), {})[str(token)] = int( + frequency + ) + + revision_pairs = sorted(revision_by_document.items()) + index_rows = ( + ( + await db.execute( + select(DocumentMapUnitIndex).where( + tuple_( + DocumentMapUnitIndex.document_id, + DocumentMapUnitIndex.job_result_id, + ).in_(revision_pairs) + ) + ) + ) + .scalars() + .all() + ) + compatible_indexes = [ + row for row in index_rows if row.format_version == MAP_UNIT_INDEX_FORMAT_VERSION + ] + if len(compatible_indexes) != len(revision_pairs): + return None + average_idf_path = combine_average_idf( + [(row.average_idf_path, row.unit_count) for row in compatible_indexes] + ) + average_idf_content = combine_average_idf( + [(row.average_idf_content, row.unit_count) for row in compatible_indexes] + ) + path_document_count = sum(row.path_document_count or 0 for row in compatible_indexes) + path_total_length = sum(row.path_total_length or 0 for row in compatible_indexes) + content_document_count = sum( + row.content_document_count or 0 for row in compatible_indexes + ) + content_total_length = sum(row.content_total_length or 0 for row in compatible_indexes) + + unit_dicts = [ + { + "map_unit_id": row.id, + "path_length": row.path_token_count, + "content_length": row.content_token_count, + } + for row in unit_rows + ] + path_stats = build_channel_bm25_stats( + unit_rows=unit_dicts, + map_unit_id_field="map_unit_id", + length_field="path_length", + channel="path", + query_tokens=query_tokens, + frequencies=frequencies, + average_idf=average_idf_path, + document_count_override=path_document_count or None, + total_length_override=path_total_length or None, + ) + content_stats = build_channel_bm25_stats( + unit_rows=unit_dicts, + map_unit_id_field="map_unit_id", + length_field="content_length", + channel="content", + query_tokens=query_tokens, + frequencies=frequencies, + average_idf=average_idf_content, + document_count_override=content_document_count or None, + total_length_override=content_total_length or None, + ) + score_units = [ + PersistedScoreUnit( + unit_id=row.id, + path_length=row.path_token_count, + content_length=row.content_token_count, + path_frequencies=frequencies.get((row.id, "path"), {}), + content_frequencies=frequencies.get((row.id, "content"), {}), + ) + for row in unit_rows + ] + corpus = PersistedScoreCorpus( + units=score_units, path_stats=path_stats, content_stats=content_stats + ) + scored = score_persisted_corpus_many(corpus, [query]) + unit_scores = scored.get(query, {}) + + section_scores: dict[str, float] = {} + for row in unit_rows: + score = float(unit_scores.get(row.id, 0.0) or 0.0) + if score > section_scores.get(row.section_id, 0.0): + section_scores[row.section_id] = score + return section_scores + + +def pool_scores_to_tree( + *, + children_by_parent: dict[str, list[str]], + roots: list[str], + leaf_scores: dict[str, float], +) -> dict[str, float]: + """Post-order max-pool: an internal node's score is the best of its own + leaf score (if it owns a unit) and every descendant's score.""" + pooled: dict[str, float] = {} + + def visit(section_id: str) -> float: + if section_id in pooled: + return pooled[section_id] + children = children_by_parent.get(section_id) or [] + best = leaf_scores.get(section_id, 0.0) + for child in children: + best = max(best, visit(child)) + pooled[section_id] = best + return best + + for root in roots: + visit(root) + return pooled + + +def best_score_ids(pooled: dict[str, float]) -> set[str]: + """Every section whose pooled score equals the best one — the best chain + (max-pooling gives each ancestor at least its best descendant's score).""" + best = max(pooled.values(), default=0.0) + return {section_id for section_id, score in pooled.items() if score == best} + + +_TOP_LEVEL_KEY = "" + + +@dataclass(frozen=True) +class MapNode: + """One row of a rendered outline/node_filter map. + + ``parent_section_id`` is the parent only when that parent is itself a + row of the same map; ``depth`` is the row's indent level. + """ + + section_id: str + parent_section_id: str | None + render: str + depth: int = 0 + score: float = 0.0 + + +@dataclass(frozen=True) +class FoldResult: + lines: list[str] + visible_ids: set[str] + hidden_count: int + overflow: bool + + +def hidden_placeholder_line(count: int, *, depth: int | None) -> str: + """The row left in place of a folded subtree (``depth=None``: top level).""" + if depth is None: + return f"(hidden {count} nodes at the top level — re-scope to expand)" + return " " * (depth + 1) + ( + f"(hidden {count} nodes under this section — " + "re-scope to its section_path to expand)" + ) + + +def render_map_lines( + nodes: list[MapNode], + *, + hidden_ids: set[str] | frozenset[str] = frozenset(), + hidden_counts: dict[str, int] | None = None, +) -> list[str]: + counts = hidden_counts or {} + lines: list[str] = [] + for node in nodes: + if node.section_id in hidden_ids: + continue + lines.append(node.render) + count = counts.get(node.section_id) + if count: + lines.append(hidden_placeholder_line(count, depth=node.depth)) + top_level = counts.get(_TOP_LEVEL_KEY) + if top_level: + lines.append(hidden_placeholder_line(top_level, depth=None)) + return lines + + +def _joined_chars(lines: list[str]) -> int: + return sum(len(line) for line in lines) + max(len(lines) - 1, 0) + + +def fold_map_nodes( + nodes: list[MapNode], + *, + char_budget: int, + keep_ids: set[str] | frozenset[str] = frozenset(), +) -> FoldResult: + """Hide lowest-score subtrees until the joined map lines fit ``char_budget``. + + Every id in ``keep_ids`` and all of its ancestors are never hidden; + a subtree is a removal candidate only when no kept row lies in it. + Candidates go lowest score first; ties follow the archived MAP-NAV + order (``_apply_budget_hide``): most descendants first, then deepest + first, then ``section_id``. The budget counts the placeholder rows left + behind, not just the visible rows. ``overflow`` is set when the rows + that must stay (plus their placeholders) still exceed ``char_budget``; + the caller fails the call rather than hiding a kept row or cutting text. + """ + by_id = {node.section_id: node for node in nodes} + children_by_parent: dict[str, list[str]] = {} + for node in nodes: + if node.parent_section_id in by_id: + children_by_parent.setdefault(str(node.parent_section_id), []).append( + node.section_id + ) + + keep: set[str] = set() + for kept_id in keep_ids: + current = by_id.get(kept_id) + while current is not None and current.section_id not in keep: + keep.add(current.section_id) + current = by_id.get(current.parent_section_id or "") + + subtree_cache: dict[str, list[str]] = {} + + def subtree_ids(section_id: str) -> list[str]: + cached = subtree_cache.get(section_id) + if cached is not None: + return cached + ids = [section_id] + for child_id in children_by_parent.get(section_id, []): + ids.extend(subtree_ids(child_id)) + subtree_cache[section_id] = ids + return ids + + row_chars = {node.section_id: len(node.render) + 1 for node in nodes} + visible_chars = sum(row_chars.values()) + hidden_ids: set[str] = set() + hidden_counts: dict[str, int] = {} + + def placeholder_chars() -> int: + total = 0 + for key, count in hidden_counts.items(): + depth = None if key == _TOP_LEVEL_KEY else by_id[key].depth + total += len(hidden_placeholder_line(count, depth=depth)) + 1 + return total + + candidates = sorted( + (node for node in nodes if node.section_id not in keep), + key=lambda node: ( + node.score, + -(len(subtree_ids(node.section_id)) - 1), + -node.depth, + node.section_id, + ), + ) + for candidate in candidates: + if visible_chars + placeholder_chars() - 1 <= char_budget: + break + if candidate.section_id in hidden_ids: + continue + ids = subtree_ids(candidate.section_id) + newly_hidden = [sid for sid in ids if sid not in hidden_ids] + folded = len(newly_hidden) + sum( + hidden_counts.pop(sid) for sid in ids if sid in hidden_counts + ) + hidden_ids.update(newly_hidden) + visible_chars -= sum(row_chars[sid] for sid in newly_hidden) + key = ( + str(candidate.parent_section_id) + if candidate.parent_section_id in by_id + else _TOP_LEVEL_KEY + ) + hidden_counts[key] = hidden_counts.get(key, 0) + folded + + lines = render_map_lines(nodes, hidden_ids=hidden_ids, hidden_counts=hidden_counts) + return FoldResult( + lines=lines, + visible_ids={node.section_id for node in nodes if node.section_id not in hidden_ids}, + hidden_count=sum(hidden_counts.values()), + overflow=_joined_chars(lines) > char_budget, + ) diff --git a/packages/shared-python/shared/services/retrieval/scoring/score_units.py b/packages/shared-python/shared/services/retrieval/scoring/score_units.py index 11ba11648..efbabbf6c 100644 --- a/packages/shared-python/shared/services/retrieval/scoring/score_units.py +++ b/packages/shared-python/shared/services/retrieval/scoring/score_units.py @@ -7,7 +7,6 @@ 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]: @@ -153,9 +152,6 @@ def build_score_units( 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 - ), } ) @@ -182,9 +178,6 @@ def build_score_units( 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/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py index b6647a16d..176dbf746 100644 --- a/packages/shared-python/shared/services/retrieval/search/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/search/lexical_text.py @@ -9,7 +9,7 @@ from shared.services.chunks.document_path import split_document_path from shared.services.chunks.path_segments import split_escaped_document_path -from shared.utils.text_utils import tokenize_contents_for_retrieval +from shared.services.retrieval.publication_preparation_cache import PublicationPreparationCache _SAME_AS_RE = re.compile(r"\[SAME-AS [^\]]+\]") @@ -44,8 +44,7 @@ def build_lexical_text(value: str) -> str: if not text: return "" - tokens = tokenize_contents_for_retrieval([text], stopwords=[], link_char=" ", dedupe=True) - token_text = tokens[0] if tokens else "" + token_text = " ".join(PublicationPreparationCache.tokenize(text, dedupe=True)) lexical_parts = [part for part in [text, token_text] if part] return "\n".join(lexical_parts) if lexical_parts else text @@ -136,8 +135,7 @@ def build_content_search_text( if section_summary and str(section_summary).strip(): parts.append(str(section_summary).strip()) raw = " ".join(parts) - tokens = tokenize_contents_for_retrieval([raw], stopwords=[], link_char=" ") - return tokens[0] if tokens else raw + return " ".join(PublicationPreparationCache.tokenize(raw)) def _chunk_summary_text(chunk: dict[str, Any]) -> str: @@ -179,8 +177,7 @@ def build_path_search_text( if not parts: return None raw = " ".join(parts) - tokens = tokenize_contents_for_retrieval([raw], stopwords=[], link_char=" ") - return tokens[0] if tokens else raw + return " ".join(PublicationPreparationCache.tokenize(raw)) def build_term_search_text( 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 3a96fbdd1..d0ad50a2d 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 @@ -98,6 +98,7 @@ async def _publish() -> None: {document_scope_clause} {type_clause} {signal_clause} + {section_clause} ) """ @@ -189,6 +190,36 @@ def _build_signal_clause( return clause, params +def _build_section_subtree_clause( + targets: list[tuple[str, str | None]] | None, +) -> tuple[str, dict[str, Any]]: + """Restrict units to exact section subtrees. Empty when no target names a path. + + Whole-document targets stay ``document_id`` only. A named path is that + section and everything under it (``path = X OR path LIKE 'X / %'``), not + a substring match. + """ + if not targets or not any(section_path for _document_id, section_path in targets): + return "", {} + parts: list[str] = [] + params: dict[str, Any] = {} + for index, (document_id, section_path) in enumerate(targets): + doc_key = f"_sec_doc_{index}" + params[doc_key] = document_id + if not section_path: + parts.append(f"(dmu.document_id = :{doc_key})") + continue + path_key = f"_sec_path_{index}" + like_key = f"_sec_pathlike_{index}" + params[path_key] = section_path + params[like_key] = f"{section_path} / %" + parts.append( + f"(dmu.document_id = :{doc_key} AND " + f"(ds.section_path = :{path_key} OR ds.section_path LIKE :{like_key}))" + ) + return f"AND ({' OR '.join(parts)})", params + + async def map_unit_discovery( db: AsyncSession | None, *, @@ -204,6 +235,7 @@ async def map_unit_discovery( filter_mode: str = "delete", revision_pins: Mapping[str, str] | None = None, publish_index_readiness: bool = True, + section_targets: list[tuple[str, str | None]] | None = None, **_kwargs: Any, ) -> DiscoveryResult: """Score the whole in-scope corpus via the persisted map-unit BM25 scorer.""" @@ -228,22 +260,26 @@ async def map_unit_discovery( signal_clause, signal_params = _build_signal_clause( signal_paths or [], filter_mode ) + section_clause, section_params = _build_section_subtree_clause(section_targets) params: dict[str, Any] = {"user_id": user_id, "namespace": namespace} params.update(revision_params) params.update(document_scope_params) params.update(type_params) params.update(signal_params) + params.update(section_params) is_unfiltered_scope: bool = not any( ( chunk_types, signal_paths, + section_clause, exclude_sections, exclude_document_ids, document_scope.include is not None, document_scope.exclude, ) ) + uses_section_join = bool(signal_paths or section_clause) cte = _SCOPED_UNITS_CTE.format( revision_join=revision_join, @@ -251,6 +287,7 @@ async def map_unit_discovery( document_scope_clause=document_scope_clause, type_clause=type_clause, signal_clause=signal_clause, + section_clause=section_clause, ) unit_statement = cte + "SELECT * FROM scoped_units" if is_unfiltered_scope: @@ -288,7 +325,7 @@ async def map_unit_discovery( len(unit_rows), ) - frequency_scope_cte = cte if signal_paths else _SCOPED_UNIT_IDS_CTE.format( + frequency_scope_cte = cte if uses_section_join else _SCOPED_UNIT_IDS_CTE.format( revision_join=revision_join, revision_clause=revision_clause, document_scope_clause=document_scope_clause, @@ -357,7 +394,7 @@ async def map_unit_discovery( index_statement = ( ( cte - if signal_paths + if uses_section_join else _SCOPED_UNIT_IDS_CTE.format( revision_join=revision_join, revision_clause=revision_clause, @@ -484,7 +521,7 @@ async def map_unit_discovery( # undercount check used for unfiltered token projection. has_index_unit_count_mismatch = ( indexed_unit_count < len(unit_rows) - if is_unfiltered_scope or type_clause or signal_paths or exclude_sections + if is_unfiltered_scope or type_clause or signal_paths or section_clause or exclude_sections else indexed_unit_count != len(unit_rows) ) is_index_format_incompatible = any( @@ -765,7 +802,8 @@ async def _hydrate_winning_units( "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 " + "jr.job_id, " + "jr.document_metadata ->> 'result_raw_prefix' AS result_raw_prefix " "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 " diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index 47f27d5a8..69c581fb0 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -190,6 +190,10 @@ async def load_all_scoped_chunks( 'chunk_metadata': chunk.chunk_metadata or {}, 'job_result_id': chunk.job_result_id, 'job_id': job_result.job_id if job_result else None, + 'result_raw_prefix': ( + (job_result.document_metadata or {}).get('result_raw_prefix') + if job_result else None + ), 'sort_order': chunk.sort_order, }) return rows diff --git a/packages/shared-python/shared/services/retrieval/search/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py index af546ae8e..e60003741 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoring.py +++ b/packages/shared-python/shared/services/retrieval/search/scoring.py @@ -2,47 +2,12 @@ from typing import Any -from shared.services.retrieval.settings import RRF_K - def get_row_path(row: dict[str, Any]) -> str: """Extract the canonical path from a row for deduplication.""" return str(row.get('section_path') or row.get('source_chunk_path') or '') -def merge_channels_rrf( - channels: list[list[dict[str, Any]]], - weights: list[float], - top_k: int, - k: int = RRF_K, -) -> list[dict[str, Any]]: - """Reciprocal Rank Fusion across multiple retrieval channels.""" - score_dict: dict[str, float] = {} - row_by_chunk_id: dict[str, dict[str, Any]] = {} - - for channel_idx, channel_rows in enumerate(channels): - weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - seen_chunk_ids: set[str] = set() - unique_rank = 0 - for row in channel_rows: - chunk_id = str(row.get('chunk_id') or '') - if not chunk_id or chunk_id in seen_chunk_ids: - continue - seen_chunk_ids.add(chunk_id) - rrf_score = weight / (k + unique_rank + 1) - score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score - if chunk_id not in row_by_chunk_id: - row_by_chunk_id[chunk_id] = row - unique_rank += 1 - - ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True) - results: list[dict[str, Any]] = [] - for chunk_id, fused_score in ranked[:top_k]: - row = row_by_chunk_id[chunk_id] - results.append(dict(row, score=round(fused_score, 6))) - return results - - def normalize_row_scores( rows: list[dict[str, Any]], *, diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py index 5d359247a..4768177a1 100644 --- a/packages/shared-python/shared/services/retrieval/settings.py +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -1,9 +1,6 @@ from __future__ import annotations -CHANNEL_WEIGHT_PATH = 1.0 -CHANNEL_WEIGHT_CONTENT = 2.0 INTERNAL_RECALL_K_MULTIPLIER = 2 -RRF_K = 60 DEFAULT_TOP_K = 10 # Final evidence / tool-observation text budget (characters). Used by diff --git a/packages/shared-python/shared/services/storage/page_pdf_crop.py b/packages/shared-python/shared/services/storage/page_pdf_crop.py index c3de8f514..9024f574e 100644 --- a/packages/shared-python/shared/services/storage/page_pdf_crop.py +++ b/packages/shared-python/shared/services/storage/page_pdf_crop.py @@ -16,6 +16,7 @@ def crop_source_pdf_pages( *, job_id: str, + raw_prefix: str | None = None, pages: list[int], storage: JobResultStorage | None = None, expires_in: int = 3600, @@ -27,22 +28,36 @@ def crop_source_pdf_pages( result_storage = storage or JobResultStorage() cache_ref = _cache_ref(normalized_pages) - if result_storage.verify_raw_exists(job_id=job_id, relative_path=cache_ref): - return result_storage.generate_artifact_url( + if _verify_raw_exists( + result_storage, + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=cache_ref, + ): + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=cache_ref, expires_in=expires_in, ) - if not result_storage.verify_raw_exists(job_id=job_id, relative_path=_SOURCE_PDF_REF): + if not _verify_raw_exists( + result_storage, + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=_SOURCE_PDF_REF, + ): logger.warning("[page_pdf_crop] source.pdf missing for job_id={}", job_id) return None work_dir = temp_dir or tempfile.gettempdir() try: with tempfile.TemporaryDirectory(dir=work_dir) as local_dir: - source_path = result_storage.download_raw_to_temp( + source_path = _download_raw_to_temp( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, relative_path=_SOURCE_PDF_REF, suffix=".pdf", temp_dir=local_dir, @@ -53,8 +68,10 @@ def crop_source_pdf_pages( output_path=cropped_path, pages=normalized_pages, ) - result_storage.upload_raw_file( + _upload_raw_file( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, relative_path=cache_ref, local_file_path=cropped_path, ) @@ -67,13 +84,101 @@ def crop_source_pdf_pages( ) return None - return result_storage.generate_artifact_url( + return _generate_artifact_url( + result_storage, job_id=job_id, + raw_prefix=raw_prefix, artifact_ref=cache_ref, expires_in=expires_in, ) +def _verify_raw_exists( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, +) -> bool: + if raw_prefix is None: + return storage.verify_raw_exists(job_id=job_id, relative_path=relative_path) + return storage.verify_raw_exists( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + ) + + +def _generate_artifact_url( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + artifact_ref: str, + expires_in: int, +) -> str | None: + if raw_prefix is None: + return storage.generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + return storage.generate_artifact_url( + job_id=job_id, + raw_prefix=raw_prefix, + artifact_ref=artifact_ref, + expires_in=expires_in, + ) + + +def _download_raw_to_temp( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, + suffix: str, + temp_dir: str, +) -> str: + if raw_prefix is None: + return storage.download_raw_to_temp( + job_id=job_id, + relative_path=relative_path, + suffix=suffix, + temp_dir=temp_dir, + ) + return storage.download_raw_to_temp( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + suffix=suffix, + temp_dir=temp_dir, + ) + + +def _upload_raw_file( + storage: JobResultStorage, + *, + job_id: str, + raw_prefix: str | None, + relative_path: str, + local_file_path: str, +) -> None: + if raw_prefix is None: + storage.upload_raw_file( + job_id=job_id, + relative_path=relative_path, + local_file_path=local_file_path, + ) + return + storage.upload_raw_file( + job_id=job_id, + raw_prefix=raw_prefix, + relative_path=relative_path, + local_file_path=local_file_path, + ) + + def _normalize_pages(pages: list[int]) -> list[int]: normalized: set[int] = set() for page in pages: diff --git a/packages/shared-python/shared/services/storage/raw_prefix_arguments.py b/packages/shared-python/shared/services/storage/raw_prefix_arguments.py new file mode 100644 index 000000000..7766056c4 --- /dev/null +++ b/packages/shared-python/shared/services/storage/raw_prefix_arguments.py @@ -0,0 +1,7 @@ +"""Optional artifact location arguments for legacy and canonical result readers.""" + +from typing import TypedDict + + +class RawPrefixArguments(TypedDict, total=False): + raw_prefix: str diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index c68febab4..8583e19ee 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -42,16 +42,38 @@ def upload( raise NotImplementedError def generate_artifact_url( - self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 + self, + *, + job_id: str, + artifact_ref: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: raise NotImplementedError def generate_raw_file_url( - self, *, job_id: str, relative_path: str, expires_in: int = 3600 + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: raise NotImplementedError - def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: + def verify_raw_exists( + self, *, job_id: str, relative_path: str, raw_prefix: str | None = None + ) -> bool: + raise NotImplementedError + + def upload_raw_file( + self, + *, + job_id: str, + relative_path: str, + local_file_path: str, + raw_prefix: str | None = None, + ) -> None: raise NotImplementedError def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: @@ -62,6 +84,7 @@ def download_raw_to_temp( *, job_id: str, relative_path: str, + raw_prefix: str | None = None, suffix: str, temp_dir: str, ) -> str: @@ -92,11 +115,20 @@ def build_zip_key(self, *, job_id: str) -> str: def build_raw_prefix(self, *, job_id: str) -> str: return self._job_file_storage.build_result_raw_prefix(job_id=job_id) - def build_raw_key(self, *, job_id: str, relative_path: str) -> str: + def build_raw_key( + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + ) -> str: normalized = self._normalize_raw_relative_path(relative_path) if not normalized: raise ValueError(f"Invalid result raw artifact path: {relative_path}") - return f"{self.build_raw_prefix(job_id=job_id)}{normalized}" + prefix = self._normalize_raw_prefix(raw_prefix) + if prefix is None: + prefix = self.build_raw_prefix(job_id=job_id) + return f"{prefix}{normalized}" def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: normalized = self._normalize_raw_relative_path(artifact_ref) @@ -205,16 +237,35 @@ def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | Non expires_in=expires_in, )["download_url"] - def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: - key = self.build_raw_key(job_id=job_id, relative_path=relative_path) + def verify_raw_exists( + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + ) -> bool: + key = self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ) result = self._job_file_storage.verify_exists(key, bucket=self.results_bucket) return bool(result.get("exists")) def generate_raw_file_url( - self, *, job_id: str, relative_path: str, expires_in: int = 3600 + self, + *, + job_id: str, + relative_path: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: return self.generate_url( - storage_key=self.build_raw_key(job_id=job_id, relative_path=relative_path), + storage_key=self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ), expires_in=expires_in, ) @@ -223,11 +274,16 @@ def upload_raw_file( *, job_id: str, relative_path: str, + raw_prefix: str | None = None, local_file_path: str, ) -> None: self._job_file_storage.upload_local_file( local_file_path, - self.build_raw_key(job_id=job_id, relative_path=relative_path), + self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ), bucket=self.results_bucket, ) @@ -236,24 +292,38 @@ def download_raw_to_temp( *, job_id: str, relative_path: str, + raw_prefix: str | None = None, suffix: str, temp_dir: str, ) -> str: return self._job_file_storage.download_to_temp( - self.build_raw_key(job_id=job_id, relative_path=relative_path), + self.build_raw_key( + job_id=job_id, + relative_path=relative_path, + raw_prefix=raw_prefix, + ), suffix=suffix, temp_dir=temp_dir, bucket=self.results_bucket, ) def generate_artifact_url( - self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 + self, + *, + job_id: str, + artifact_ref: str, + raw_prefix: str | None = None, + expires_in: int = 3600, ) -> str | None: normalized_ref = self.normalize_artifact_ref(artifact_ref) if not normalized_ref: return None return self.generate_url( - storage_key=self.build_raw_key(job_id=job_id, relative_path=normalized_ref), + storage_key=self.build_raw_key( + job_id=job_id, + relative_path=normalized_ref, + raw_prefix=raw_prefix, + ), expires_in=expires_in, ) @@ -284,6 +354,21 @@ def _normalize_raw_relative_path(self, relative_path: str | None) -> str | None: return None return "/".join(parts) + def _normalize_raw_prefix(self, raw_prefix: str | None) -> str | None: + if raw_prefix is None: + return None + normalized = str(raw_prefix).strip().replace("\\", "/").lstrip("/") + if not normalized: + raise ValueError("Result raw prefix must not be empty") + if not normalized.startswith("results/"): + raise ValueError("Result raw prefix must be under results/") + parts = normalized.rstrip("/").split("/") + if not parts or any(not part or part in {".", ".."} for part in parts): + raise ValueError("Result raw prefix contains an invalid path component") + if any(self._is_excluded_dir(part) for part in parts): + raise ValueError("Result raw prefix contains an excluded directory") + return "/".join(parts) + "/" + def _normalize_artifact_refs(self, artifact_refs: set[str] | None) -> set[str] | None: if artifact_refs is None: return None diff --git a/packages/shared-python/shared/tests/test_agent_explore_harness.py b/packages/shared-python/shared/tests/test_agent_explore_harness.py index 113247db6..553793b45 100644 --- a/packages/shared-python/shared/tests/test_agent_explore_harness.py +++ b/packages/shared-python/shared/tests/test_agent_explore_harness.py @@ -72,12 +72,15 @@ def test_build_wire_tool_name_map_round_trips_to_canonical() -> None: 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" + assert ( + tool_message_content(result, tool_name="corpus.grep", 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) + content = tool_message_content(result, tool_name="corpus.grep", 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 @@ -85,12 +88,18 @@ def test_tool_message_content_caps_long_text_with_note() -> None: 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" + assert ( + tool_message_content(result, tool_name="corpus.grep", 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)" + assert ( + tool_message_content(result, tool_name="corpus.grep", max_chars=100) + == "(empty result)" + ) def test_model_accepts_images_only_when_name_contains_vision() -> None: @@ -329,7 +338,8 @@ def test_agent_explore_keeps_inventory_tool_for_explicit_inventory_requests() -> wire_names = {tool["function"]["name"] for tool in tools} assert "corpus_list_documents" in wire_names assert name_map["corpus_list_documents"] == "corpus.list_documents" - assert "only if the user explicitly asks to list or inventory" in LOOP_CONTRACT_SUFFIX + assert "only if the user explicitly" in LOOP_CONTRACT_SUFFIX + assert "asks to list or inventory the corpus's documents" in LOOP_CONTRACT_SUFFIX # -------------------------------------------------------------------------- diff --git a/packages/shared-python/shared/tests/test_agent_tools_contract.py b/packages/shared-python/shared/tests/test_agent_tools_contract.py new file mode 100644 index 000000000..d334ff2d3 --- /dev/null +++ b/packages/shared-python/shared/tests/test_agent_tools_contract.py @@ -0,0 +1,258 @@ +"""Schema contract: closed objects, field descriptions, unknown/old names error.""" + +from __future__ import annotations + +import os +from contextlib import asynccontextmanager + +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.config import FINISH_TOOL_SCHEMA +from shared.services.retrieval.agent_tools import REGISTRY, ToolContext, load_corpus_schema_text +from shared.services.retrieval.agent_tools.snippet import build_row, format_row + + +@asynccontextmanager +async def _unused_db_factory(): + raise AssertionError("schema validation must not open a database") + yield # pragma: no cover + + +def _ctx() -> ToolContext: + return ToolContext( + db=None, # type: ignore[arg-type] + user_id="user_contract", + namespace="default", + db_factory=_unused_db_factory, + ) + + +def _walk_objects(schema: object, path: str) -> list[tuple[str, dict]]: + found: list[tuple[str, dict]] = [] + if not isinstance(schema, dict): + return found + if schema.get("type") == "object" or "properties" in schema: + found.append((path, schema)) + for name, child in (schema.get("properties") or {}).items(): + found.extend(_walk_objects(child, f"{path}.{name}")) + if "items" in schema: + found.extend(_walk_objects(schema["items"], f"{path}[]")) + for key in ("anyOf", "oneOf", "allOf"): + for index, child in enumerate(schema.get(key) or []): + found.extend(_walk_objects(child, f"{path}/{key}[{index}]")) + return found + + +def test_every_schema_object_closes_unknown_keys() -> None: + for spec in REGISTRY.all(): + for path, obj in _walk_objects(spec.json_schema, spec.name): + assert obj.get("additionalProperties") is False, path + for path, obj in _walk_objects(FINISH_TOOL_SCHEMA, "finish"): + assert obj.get("additionalProperties") is False, path + + +def test_every_schema_property_has_a_description() -> None: + for spec in REGISTRY.all(): + for path, obj in _walk_objects(spec.json_schema, spec.name): + for name, child in (obj.get("properties") or {}).items(): + if not isinstance(child, dict): + continue + if child.get("type") == "object" or "properties" in child: + continue + assert child.get("description"), f"{path}.{name}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_name", "args", "legal_names"), + [ + ( + "corpus.grep", + {"pattern": "x", "document_ids": ["doc_a"]}, + ["chunk_types", "context_chars", "limit", "pattern", "patterns", "scope"], + ), + ( + "corpus.grep", + {"pattern": "x", "max_results": 3}, + ["chunk_types", "context_chars", "limit", "pattern", "patterns", "scope"], + ), + ( + "corpus.recall", + {"query": "x", "top_k": 5}, + ["chunk_types", "limit", "query", "scope"], + ), + ( + "corpus.recall", + {"query": "x", "channels": ["term"]}, + ["chunk_types", "limit", "query", "scope"], + ), + ( + "corpus.outline", + {"scope": [{"document_id": "doc_a"}], "document_id": "doc_a"}, + ["depth", "scope"], + ), + ( + "corpus.outline", + {"scope": [{"document_id": "doc_a"}], "path_prefix": "A"}, + ["depth", "scope"], + ), + ( + "corpus.assets", + {"document_ids": ["doc_a"]}, + ["host_of", "query", "scope", "type"], + ), + ( + "corpus.read", + { + "refs": [{"document_id": "doc_a", "section_path": "A / B"}], + "section_paths": ["A / B"], + }, + ["include_assets", "mode", "refs", "resolve_same_as"], + ), + ( + "corpus.node_filter", + { + "scope": [{"document_id": "doc_a"}], + "predicates": [{"field": "path", "terms": ["A"]}], + "document_ids": ["doc_a"], + }, + ["chunk_types", "predicates", "scope"], + ), + ], +) +async def test_unknown_or_old_names_fail_with_legal_parameter_list( + tool_name: str, args: dict, legal_names: list[str] +) -> None: + result = await REGISTRY.dispatch(tool_name, _ctx(), args) + assert result.error is not None + assert "unknown argument" in result.error + for name in legal_names: + assert name in result.error + + +@pytest.mark.asyncio +async def test_nested_unknown_key_names_the_path() -> None: + result = await REGISTRY.dispatch( + "corpus.grep", + _ctx(), + {"pattern": "x", "scope": [{"document_id": "doc_a", "bogus": 1}]}, + ) + assert result.error is not None + assert "scope[0]" in result.error + assert "document_id" in result.error + assert "section_path" in result.error + + +def test_old_parameter_names_are_absent_from_docs_and_descriptions() -> None: + schema_text = load_corpus_schema_text() + for stale in ("document_ids", "path_prefix", "max_results"): + assert stale not in schema_text + for spec in REGISTRY.all(): + assert stale not in spec.description + assert stale not in spec.json_schema.get("properties", {}) + + +_ROW_KEYS = { + "kind", + "title", + "document_id", + "section_path", + "chunk_id", + "summary", + "snippet", + "score", + "is_hit", + "depth", + "hosted", +} + + +def test_format_row_is_the_shared_search_line() -> None: + row = build_row( + kind="text", + document_id="doc_a", + section_path="guide.pdf / Intro", + title="guide.pdf", + snippet="hello", + ) + assert set(row) == _ROW_KEYS + line = format_row(row) + assert line.startswith("- [text] guide.pdf | document_id=doc_a section_path=") + assert "snippet:" in line + + +@pytest.mark.asyncio +async def test_chunk_types_enum_rejects_old_names() -> None: + result = await REGISTRY.dispatch( + "corpus.grep", + _ctx(), + {"pattern": "x", "chunk_types": ["body"]}, + ) + assert result.error is not None + assert "chunk_types" in result.error + assert "unknown argument" not in result.error + + +@pytest.mark.asyncio +async def test_limit_below_minimum_is_rejected() -> None: + result = await REGISTRY.dispatch( + "corpus.grep", + _ctx(), + {"pattern": "x", "limit": 0}, + ) + assert result.error is not None + assert "limit" in result.error + + +@pytest.mark.asyncio +async def test_recall_blank_query_is_rejected() -> None: + result = await REGISTRY.dispatch("corpus.recall", _ctx(), {"query": " "}) + assert result.error == "recall requires query" + + +@pytest.mark.asyncio +async def test_read_ref_requires_exactly_one_address() -> None: + both = await REGISTRY.dispatch( + "corpus.read", + _ctx(), + { + "refs": [ + { + "document_id": "doc_a", + "section_path": "guide.pdf / Intro", + "chunk_id": "chunk_a", + } + ] + }, + ) + assert both.error is not None + assert "section_path" in both.error + assert "chunk_id" in both.error + + neither = await REGISTRY.dispatch( + "corpus.read", + _ctx(), + {"refs": [{"document_id": "doc_a"}]}, + ) + assert neither.error is not None + assert "section_path" in neither.error or "chunk_id" in neither.error + + +def test_chunk_types_and_bounds_are_in_schema() -> None: + grep = REGISTRY.get("corpus.grep") + recall = REGISTRY.get("corpus.recall") + outline = REGISTRY.get("corpus.outline") + assert grep is not None and recall is not None and outline is not None + for spec in (grep, recall): + items = spec.json_schema["properties"]["chunk_types"]["items"] + assert items["enum"] == ["text", "page", "image", "table"] + assert spec.json_schema["properties"]["limit"]["minimum"] == 1 + assert grep.json_schema["properties"]["context_chars"]["minimum"] == 1 + assert outline.json_schema["properties"]["depth"]["minimum"] == 0 diff --git a/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py b/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py index d376e4306..a74718c93 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py +++ b/packages/shared-python/shared/tests/test_agent_tools_discovery_read_contract.py @@ -65,6 +65,7 @@ def all(self) -> list[tuple[object, ...]]: {"summary": "dose table"}, REV_ID, JOB_ID, + None, PATH_ROOT, FILE_NAME, 1, @@ -72,12 +73,20 @@ def all(self) -> list[tuple[object, ...]]: ] +class _EmptyHostRows: + def all(self) -> list[tuple[object, ...]]: + return [] + + class _RecordingDb: - def __init__(self, result: object) -> None: - self.result = result + def __init__(self, results: object) -> None: + self.results = list(results) if isinstance(results, list) else [results] + self._index = 0 async def execute(self, statement): # noqa: ANN001 - return self.result + result = self.results[min(self._index, len(self.results) - 1)] + self._index += 1 + return result @asynccontextmanager @@ -199,21 +208,21 @@ def _read_kwargs(refs: list[dict[str, str]]) -> dict: def _section_paths_from_outline_text(text: str) -> list[str]: - return re.findall(r"\| section_path=(.+) \(chunks=", text) + return re.findall(r"section_path=(.+)$", text, re.MULTILINE) def _asset_refs_from_text(text: str) -> list[dict[str, str]]: return [ {"document_id": document_id, "chunk_id": chunk_id} for document_id, chunk_id in re.findall( - r"document_id=(\S+) chunk_id=(\S+)", text + r"document_id=(\S+) section_path=.+? chunk_id=(\S+)", text ) ] @pytest.mark.asyncio async def test_outline_visible_section_path_reads(discovery_ctx: ToolContext) -> None: - listed = await outline(discovery_ctx, {"document_id": DOC_ID}) + listed = await outline(discovery_ctx, {"scope": [{"document_id": DOC_ID}]}) assert listed.error is None paths = _section_paths_from_outline_text(listed.text) assert PATH_INTRO in paths @@ -222,33 +231,36 @@ async def test_outline_visible_section_path_reads(discovery_ctx: ToolContext) -> _read_kwargs([{"document_id": DOC_ID, "section_path": PATH_INTRO}]), ) assert result.error is None - assert result.payload["errors"] == [] + assert [entry["status"] for entry in result.payload["refs"]] == ["ok"] assert result.refs == [{"document_id": DOC_ID, "chunk_id": CHUNK_INTRO}] @pytest.mark.asyncio async def test_assets_visible_chunk_id_reads(discovery_ctx: ToolContext) -> None: - listed = await assets(discovery_ctx, {"document_ids": [DOC_ID], "type": "table"}) + listed = await assets( + discovery_ctx, {"scope": [{"document_id": DOC_ID}], "type": "table"} + ) assert listed.error is None - assert TABLE_FILE in listed.text + assert listed.payload["rows"][0]["chunk_id"] == CHUNK_TABLE + assert listed.payload["rows"][0]["kind"] == "table" refs = _asset_refs_from_text(listed.text) assert refs == [{"document_id": DOC_ID, "chunk_id": CHUNK_TABLE}] result = await read(discovery_ctx, _read_kwargs(refs)) assert result.error is None - assert result.payload["errors"] == [] + assert [entry["status"] for entry in result.payload["refs"]] == ["ok"] assert result.refs == [{"document_id": DOC_ID, "chunk_id": CHUNK_TABLE}] @pytest.mark.asyncio async def test_grep_table_visible_chunk_id_reads(discovery_ctx: ToolContext) -> None: grep_ctx = ToolContext( - db=_RecordingDb(_GrepRows()), # type: ignore[arg-type] + db=_RecordingDb([_GrepRows(), _EmptyHostRows()]), # type: ignore[arg-type] user_id=USER_ID, namespace=NAMESPACE, db_factory=_unused_db_factory, ) listed = await grep(grep_ctx, {"pattern": "30 mg"}) - match = re.search(r"\((\S+)\) chunk_id=(\S+) /", listed.text) + match = re.search(r"document_id=(\S+) section_path=.+? chunk_id=(\S+)", listed.text) assert match is not None result = await read( discovery_ctx, @@ -257,7 +269,7 @@ async def test_grep_table_visible_chunk_id_reads(discovery_ctx: ToolContext) -> ), ) assert result.error is None - assert result.payload["errors"] == [] + assert [entry["status"] for entry in result.payload["refs"]] == ["ok"] assert result.refs == [{"document_id": DOC_ID, "chunk_id": CHUNK_TABLE}] @@ -270,4 +282,6 @@ async def test_file_path_is_not_a_readable_chunk_id( _read_kwargs([{"document_id": DOC_ID, "chunk_id": TABLE_FILE}]), ) assert result.error is not None - assert f"unknown chunk_id: {TABLE_FILE}" in result.error + failed_entry = result.payload["refs"][0] + assert failed_entry["status"] == "failed" + assert f"unknown chunk_id: {TABLE_FILE}" in failed_entry["reason"] diff --git a/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py b/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py index 20639e557..845af133e 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py +++ b/packages/shared-python/shared/tests/test_agent_tools_explore_table_image.py @@ -572,37 +572,42 @@ async def test_recall_mounts_connected_table( explore_ctx: ToolContext, monkeypatch: pytest.MonkeyPatch, ) -> None: - async def _term_rows(*_args, **_kwargs): # noqa: ANN001 - return [ - { - "chunk_id": CHUNK_TEXT, - "document_id": DOC_ID, - "section_id": "sec_intro", - "section_path": PATH_INTRO, - "source_file_name": FILE_NAME, - "chunk_type": "text", - "snippet": "intro body", - "content": "intro body [tables/small.html]", - "file_path": None, - "chunk_metadata": { - "connect_to": [ - { - "target": CHUNK_SMALL, - "relation": "embeds", - "ref": "[tables/small.html]", - } - ] - }, - "job_result_id": REV_ID, - "job_id": JOB_ID, - } - ] + async def _discovery(*_args, **_kwargs): # noqa: ANN001 + return DiscoveryResult( + status="discovery_done", + payload={ + "fused_rows": [ + { + "chunk_id": CHUNK_TEXT, + "document_id": DOC_ID, + "section_id": "sec_intro", + "section_path": PATH_INTRO, + "source_file_name": FILE_NAME, + "chunk_type": "text", + "snippet": "intro body", + "content": "intro body [tables/small.html]", + "file_path": None, + "chunk_metadata": { + "connect_to": [ + { + "target": CHUNK_SMALL, + "relation": "embeds", + "ref": "[tables/small.html]", + } + ] + }, + "job_result_id": REV_ID, + "job_id": JOB_ID, + } + ] + }, + ) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall._term_channel_rows", - _term_rows, + "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + _discovery, ) - result = await recall(explore_ctx, {"query": "intro", "channels": ["term"]}) + result = await recall(explore_ctx, {"query": "intro"}) assert result.error is None assert f"mounted table chunk_id={CHUNK_SMALL}:" in result.text assert " None: - async def _term_rows(*_args, **_kwargs): # noqa: ANN001 - return [ - { - "chunk_id": CHUNK_SMALL, - "document_id": DOC_ID, - "section_id": "sec_root", - "section_path": PATH_ROOT, - "source_file_name": FILE_NAME, - "chunk_type": "table", - "snippet": "dose table summary", - "content": SMALL_HTML, - "file_path": "tables/small.html", - "chunk_metadata": {"summary": "dose table summary"}, - "job_result_id": REV_ID, - "job_id": JOB_ID, - } - ] + async def _discovery(*_args, **_kwargs): # noqa: ANN001 + return DiscoveryResult( + status="discovery_done", + payload={ + "fused_rows": [ + { + "chunk_id": CHUNK_SMALL, + "document_id": DOC_ID, + "section_id": "sec_root", + "section_path": PATH_ROOT, + "source_file_name": FILE_NAME, + "chunk_type": "table", + "snippet": "dose table summary", + "content": SMALL_HTML, + "file_path": "tables/small.html", + "chunk_metadata": {"summary": "dose table summary"}, + "job_result_id": REV_ID, + "job_id": JOB_ID, + } + ] + }, + ) monkeypatch.setattr( - "shared.services.retrieval.agent_tools.tools.recall._term_channel_rows", - _term_rows, + "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + _discovery, ) - result = await recall(explore_ctx, {"query": "dose", "channels": ["term"]}) + result = await recall(explore_ctx, {"query": "dose"}) assert result.error is None assert f"chunk_id={CHUNK_SMALL}" in result.text assert " list[tuple[object, ...]]: }, REV_ID, JOB_ID, + None, PATH_INTRO, FILE_NAME, 1, diff --git a/packages/shared-python/shared/tests/test_agent_tools_grep.py b/packages/shared-python/shared/tests/test_agent_tools_grep.py index b621188f4..67437f1af 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_grep.py +++ b/packages/shared-python/shared/tests/test_agent_tools_grep.py @@ -4,6 +4,8 @@ import os from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Any os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") @@ -17,7 +19,7 @@ from sqlalchemy.dialects import postgresql from shared.services.retrieval.agent_tools.registry import REGISTRY, ToolContext -from shared.services.retrieval.agent_tools.snippet import format_search_hit_line +from shared.services.retrieval.agent_tools.snippet import build_row, format_row from shared.services.retrieval.agent_tools.tools.grep import ( _term_search, _terms_from_args, @@ -44,6 +46,7 @@ def _text_row( None, "jr_a", "job_a", + None, section_path, source_file_name, total, @@ -91,22 +94,56 @@ def all(self) -> list[tuple[object, ...]]: return [] -class _RecordingDb: - def __init__(self, result: object) -> None: - self.result = result +class _ScalarsResult: + """Fakes ``(await db.execute(...)).scalars().all()`` for ORM lookups + (``resolve_scope``'s ``Document`` query).""" + + def __init__(self, items: list[Any]) -> None: + self._items = items + + def scalars(self) -> "_ScalarsResult": + return self + + def all(self) -> list[Any]: + return self._items + + +class _SequencedDb: + """Replays one canned result per ``execute()`` call, in order. + + ``resolve_scope`` (when a call passes ``scope``) issues its own + ``Document`` lookup before grep's own rows query — this fake replaces + the single-result ``_RecordingDb`` so each call gets the result meant + for it. + """ + + def __init__(self, results: list[Any]) -> None: + self.results = list(results) self.execute_count = 0 - self.statement = None + self.statements: list[Any] = [] - async def execute(self, statement): # noqa: ANN001 + async def execute(self, statement: Any): # noqa: ANN001 self.execute_count += 1 - if self.statement is None: - self.statement = statement - return self.result + self.statements.append(statement) + return self.results[self.execute_count - 1] + + @property + def statement(self) -> Any: + return self.statements[0] if self.statements else None + + @property + def rows_statement(self) -> Any: + """The final (rows) statement — last one issued.""" + return self.statements[-1] if self.statements else None + + +def _scoped_document(document_id: str, job_result_id: str = "jr_a") -> SimpleNamespace: + return SimpleNamespace(document_id=document_id, current_job_result_id=job_result_id) @pytest.mark.asyncio async def test_grep_runs_one_scoped_query_with_exact_total() -> None: - db = _RecordingDb(_RowsResult()) + db = _SequencedDb([_RowsResult()]) @asynccontextmanager async def rows_factory(): @@ -119,19 +156,19 @@ async def rows_factory(): namespace="default", db_factory=rows_factory, ) - result = await grep(ctx, {"pattern": "HFrEF", "document_ids": ["doc_a"]}) + result = await grep(ctx, {"pattern": "HFrEF"}) assert result.error is None - assert result.payload["total_matches"] == 2 - assert [row["chunk_id"] for row in result.payload["results"]] == ["chunk_a", "chunk_b"] + assert result.payload["details"]["total_matches"] == 2 + assert [row["document_id"] for row in result.payload["rows"]] == ["doc_a", "doc_a"] + assert [row["chunk_id"] for row in result.payload["rows"]] == [None, None] assert result.refs == [ {"document_id": "doc_a", "chunk_id": "chunk_a"}, {"document_id": "doc_a", "chunk_id": "chunk_b"}, ] assert db.execute_count == 1 - assert db.statement is not None sql = str( - db.statement.compile( + db.rows_statement.compile( dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True}, ) @@ -141,7 +178,6 @@ async def rows_factory(): assert "term_search_text" in matched_sql assert "document_chunks.user_id = 'user_grep'" in matched_sql assert "document_chunks.namespace = 'default'" in matched_sql - assert "document_chunks.document_id IN ('doc_a')" in matched_sql assert "lower(" in matched_sql.lower() assert "coalesce" in matched_sql.lower() assert " like " in matched_sql.lower() @@ -149,13 +185,68 @@ async def rows_factory(): assert "~*" not in matched_sql assert "%hfref%" in matched_sql.lower() assert result.text.startswith("total_matches=2 returned=2") - assert "- [text] guide.pdf (doc_a) / guide.pdf / Intro:" in result.text + assert "- [text] guide.pdf | document_id=doc_a section_path=guide.pdf / Intro" in result.text assert "chunk_id=" not in result.text +@pytest.mark.asyncio +async def test_grep_scope_narrows_to_document_and_subtree() -> None: + db = _SequencedDb([_ScalarsResult([_scoped_document("doc_a")]), _RowsResult()]) + + @asynccontextmanager + async def unused_factory(): + raise AssertionError("grep must not open a second database connection") + yield # pragma: no cover + + ctx = ToolContext( + db=db, # type: ignore[arg-type] + user_id="user_grep", + namespace="default", + db_factory=unused_factory, + ) + result = await grep( + ctx, {"pattern": "HFrEF", "scope": [{"document_id": "doc_a"}]} + ) + + assert result.error is None + assert db.execute_count == 2 + sql = str( + db.rows_statement.compile( + dialect=postgresql.dialect(), + compile_kwargs={"literal_binds": True}, + ) + ) + matched_sql = sql.split(")\n SELECT", 1)[0] + assert "document_chunks.document_id = 'doc_a'" in matched_sql + + +@pytest.mark.asyncio +async def test_grep_unknown_scope_document_fails_the_call() -> None: + db = _SequencedDb([_ScalarsResult([])]) + + @asynccontextmanager + async def unused_factory(): + raise AssertionError("grep must not open a second database connection") + yield # pragma: no cover + + ctx = ToolContext( + db=db, # type: ignore[arg-type] + user_id="user_grep", + namespace="default", + db_factory=unused_factory, + ) + result = await grep( + ctx, {"pattern": "HFrEF", "scope": [{"document_id": "doc_missing"}]} + ) + + assert result.error is not None + assert "unknown document_id: doc_missing" in result.error + assert db.execute_count == 1 + + @pytest.mark.asyncio async def test_grep_exact_total_survives_row_limit() -> None: - db = _RecordingDb(_LimitedRowsResult()) + db = _SequencedDb([_LimitedRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -168,16 +259,16 @@ async def unused_factory(): namespace="default", db_factory=unused_factory, ) - result = await grep(ctx, {"pattern": "HFrEF", "max_results": 1}) + result = await grep(ctx, {"pattern": "HFrEF", "limit": 1}) - assert result.payload["total_matches"] == 2 - assert len(result.payload["results"]) == 1 + assert result.payload["details"]["total_matches"] == 2 + assert len(result.payload["rows"]) == 1 assert result.text.startswith("total_matches=2 returned=1") @pytest.mark.asyncio async def test_grep_empty_result_reports_zero_total() -> None: - db = _RecordingDb(_EmptyRowsResult()) + db = _SequencedDb([_EmptyRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -192,21 +283,20 @@ async def unused_factory(): ) result = await grep(ctx, {"pattern": "missing"}) - assert result.payload == {"total_matches": 0, "results": []} + assert result.payload["rows"] == [] + assert result.payload["details"]["total_matches"] == 0 assert result.refs == [] assert result.text == "total_matches=0 returned=0" @pytest.mark.asyncio async def test_grep_requires_pattern_or_patterns() -> None: - # Error path returns before any db access, so a never-entered factory is - # enough here — mirrors the other tests' db_factory shape for consistency. @asynccontextmanager async def rows_factory(): - yield _RecordingDb(_RowsResult()) + yield _SequencedDb([_RowsResult()]) ctx = ToolContext( - db=_RecordingDb(_RowsResult()), # type: ignore[arg-type] + db=_SequencedDb([_RowsResult()]), # type: ignore[arg-type] user_id="user_grep", namespace="default", db_factory=rows_factory, @@ -227,6 +317,11 @@ def test_grep_schema_requires_pattern_or_patterns() -> None: {"required": ["patterns"]}, ] assert "is_regex" not in spec.json_schema["properties"] + assert spec.json_schema["additionalProperties"] is False + assert "document_ids" not in spec.json_schema["properties"] + assert "max_results" not in spec.json_schema["properties"] + assert "limit" in spec.json_schema["properties"] + assert "scope" in spec.json_schema["properties"] def test_terms_from_args_merges_pattern_and_patterns() -> None: @@ -293,6 +388,7 @@ def all(self) -> list[tuple[object, ...]]: {"summary": "dose table summary 30 mg"}, "jr_a", "job_a", + None, "guide.pdf / Root", "guide.pdf", 1, @@ -304,7 +400,7 @@ def all(self) -> list[tuple[object, ...]]: async def test_grep_table_hit_text_includes_chunk_id( monkeypatch: pytest.MonkeyPatch, ) -> None: - db = _RecordingDb(_TableRowsResult()) + db = _SequencedDb([_TableRowsResult(), _EmptyRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -325,7 +421,8 @@ async def unused_factory(): assert result.error is None assert ( - "- [table] guide.pdf (doc_a) chunk_id=chunk_table / guide.pdf / Root:" + "- [table] guide.pdf | document_id=doc_a section_path=guide.pdf / Root " + "(no host section) chunk_id=chunk_table" in result.text ) assert "" in result.text @@ -342,7 +439,7 @@ async def test_grep_table_download_failure_does_not_fail_whole_call( warning note instead of raising.""" from shared.services.retrieval.hydration.table_grid import TableDownloadError - db = _RecordingDb(_TableRowsResult()) + db = _SequencedDb([_TableRowsResult(), _EmptyRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -365,7 +462,7 @@ def _raise_download_error(_row): result = await grep(ctx, {"pattern": "30 mg"}) assert result.error is None - assert result.payload["total_matches"] == 1 + assert result.payload["details"]["total_matches"] == 1 assert "table unavailable" in result.text assert "download failed" in result.text @@ -374,7 +471,7 @@ def _raise_download_error(_row): async def test_grep_matches_table_via_term_search_text_not_content_path( monkeypatch: pytest.MonkeyPatch, ) -> None: - db = _RecordingDb(_TableRowsResult()) + db = _SequencedDb([_TableRowsResult(), _EmptyRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -394,44 +491,55 @@ async def unused_factory(): result = await grep(ctx, {"pattern": "dose table summary"}) assert result.error is None - assert result.payload["total_matches"] == 1 - assert result.payload["results"][0]["chunk_id"] == "chunk_table" + assert result.payload["details"]["total_matches"] == 1 + assert result.payload["rows"][0]["chunk_id"] == "chunk_table" assert "dose table summary" in result.text -def test_format_search_hit_line_body_omits_chunk_id() -> None: - assert format_search_hit_line( - source_file_name="guide.pdf", - document_id="doc_a", - section_path="guide.pdf / Intro", - snippet="alpha", - chunk_type="text", - ) == "- [text] guide.pdf (doc_a) / guide.pdf / Intro: 'alpha'" - - -def test_format_search_hit_line_omits_empty_snippet() -> None: - assert format_search_hit_line( - source_file_name="guide.pdf", - document_id="doc_a", - section_path="guide.pdf / Root", - snippet="", - chunk_type="table", - chunk_id="chunk_table", - ) == "- [table] guide.pdf (doc_a) chunk_id=chunk_table / guide.pdf / Root" - - -def test_format_search_hit_line_asset_includes_chunk_id() -> None: - assert format_search_hit_line( - source_file_name="guide.pdf", - document_id="doc_a", - section_path="guide.pdf / Root", - snippet="30 mg", - chunk_type="table", - chunk_id="chunk_table", - score=1.5, +def test_format_row_body_omits_chunk_id() -> None: + assert format_row( + build_row( + kind="text", + document_id="doc_a", + section_path="guide.pdf / Intro", + title="guide.pdf", + snippet="alpha", + ) ) == ( - "- [table] guide.pdf (doc_a) chunk_id=chunk_table / " - "guide.pdf / Root score=1.5: '30 mg'" + "- [text] guide.pdf | document_id=doc_a section_path=guide.pdf / Intro\n" + " snippet: 'alpha'" + ) + + +def test_format_row_omits_empty_snippet() -> None: + assert ( + format_row( + build_row( + kind="table", + document_id="doc_a", + section_path="guide.pdf / Root", + title="guide.pdf", + chunk_id="chunk_table", + ) + ) + == "- [table] guide.pdf | document_id=doc_a section_path=guide.pdf / Root chunk_id=chunk_table" + ) + + +def test_format_row_asset_includes_chunk_id_and_score() -> None: + assert format_row( + build_row( + kind="table", + document_id="doc_a", + section_path="guide.pdf / Root", + title="guide.pdf", + chunk_id="chunk_table", + snippet="30 mg", + score=1.5, + ) + ) == ( + "- [table] guide.pdf | document_id=doc_a section_path=guide.pdf / Root " + "chunk_id=chunk_table score=1.5\n snippet: '30 mg'" ) @@ -450,13 +558,14 @@ def all(self) -> list[tuple[object, ...]]: None, "jr_a", "job_a", + None, "guide.pdf / Root", "guide.pdf", 1, ) ] - db = _RecordingDb(_ImageRows()) + db = _SequencedDb([_ImageRows(), _EmptyRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -471,14 +580,14 @@ async def unused_factory(): ) result = await grep(ctx, {"pattern": "chart of dose"}) assert result.error is None - assert result.payload["results"][0]["chunk_id"] == "chunk_image" + assert result.payload["rows"][0]["chunk_id"] == "chunk_image" assert "chart of dose" in result.text assert "[Image:" in result.text @pytest.mark.asyncio -async def test_grep_document_ids_does_not_scan_table_cells() -> None: - db = _RecordingDb(_EmptyRowsResult()) +async def test_grep_scope_does_not_scan_table_cells() -> None: + db = _SequencedDb([_ScalarsResult([_scoped_document("doc_a")]), _EmptyRowsResult()]) @asynccontextmanager async def unused_factory(): @@ -491,7 +600,9 @@ async def unused_factory(): namespace="default", db_factory=unused_factory, ) - result = await grep(ctx, {"pattern": "30 mg", "document_ids": ["doc_a"]}) - assert db.execute_count == 1 - assert result.payload["total_matches"] == 0 + result = await grep( + ctx, {"pattern": "30 mg", "scope": [{"document_id": "doc_a"}]} + ) + assert db.execute_count == 2 + assert result.payload["details"]["total_matches"] == 0 assert "cell=" not in result.text diff --git a/packages/shared-python/shared/tests/test_agent_tools_map_scope.py b/packages/shared-python/shared/tests/test_agent_tools_map_scope.py new file mode 100644 index 000000000..1676660d7 --- /dev/null +++ b/packages/shared-python/shared/tests/test_agent_tools_map_scope.py @@ -0,0 +1,440 @@ +"""outline / node_filter / scope behavior from the tool-contract plan.""" + +from __future__ import annotations + +import os +from contextlib import asynccontextmanager +from datetime import datetime +from typing import Any + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import pytest +from sqlalchemy import create_engine, event +from sqlalchemy.orm import Session + +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 ( + MAP_TOOL_CHAR_BUDGET, + ToolContext, + ToolResult, +) +from shared.services.retrieval.agent_tools.tools.assets import assets +from shared.services.retrieval.agent_tools.tools.grep import grep +from shared.services.retrieval.agent_tools.tools.node_filter import node_filter +from shared.services.retrieval.agent_tools.tools.outline import outline +from shared.services.retrieval.agent_tools.tools.recall import recall +from shared.services.retrieval.search.map_unit_discovery import DiscoveryResult +from shared.services.retrieval.scoring.map_lighting import MapNode, fold_map_nodes + +_ROW_KEYS = { + "kind", + "title", + "document_id", + "section_path", + "chunk_id", + "summary", + "snippet", + "score", + "is_hit", + "depth", + "hosted", +} + + +def _assert_shared_search_payload(result: ToolResult) -> None: + assert set(result.payload) == {"rows", "details"} + assert isinstance(result.payload["details"], dict) + for row in result.payload["rows"]: + assert set(row) == _ROW_KEYS + +USER_ID = "user_map" +NAMESPACE = "default" +DOC_A = "doc_a" +DOC_B = "doc_b" +REV_A = "jr_a" +REV_B = "jr_b" +FILE_A = "guide.pdf" +FILE_B = "notes.pdf" +PATH_ROOT = f"{FILE_A} / Root" +PATH_OVERVIEW = f"{FILE_A} / 1 Overview" +PATH_FINDINGS = f"{FILE_A} / 1 Overview / 1.1 Findings" +PATH_DETAIL = f"{FILE_A} / 1 Overview / 1.1 Findings / Detail" +PATH_TREATMENT = f"{FILE_A} / 2 Treatment" +PATH_INTRO = f"{FILE_B} / Intro" + + +class _AsyncSessionAdapter: + def __init__(self, session: Session) -> None: + self._session = session + + async def execute(self, statement): # noqa: ANN001 + return self._session.execute(statement) + + +@asynccontextmanager +async def _unused_db_factory(): + raise AssertionError("map-scope tests should not open a second session") + yield # pragma: no cover + + +def _seed(session: Session) -> None: + now = datetime(2026, 1, 1) + session.add_all( + [ + JobResult(id=REV_A, job_id="job_a", delivery_mode="inline", created_at=now, updated_at=now), + JobResult(id=REV_B, job_id="job_b", delivery_mode="inline", created_at=now, updated_at=now), + Document( + document_id=DOC_A, + user_id=USER_ID, + namespace=NAMESPACE, + status="active", + current_job_result_id=REV_A, + source_file_name=FILE_A, + parse_track="chunk", + created_at=now, + updated_at=now, + ), + Document( + document_id=DOC_B, + user_id=USER_ID, + namespace=NAMESPACE, + status="active", + current_job_result_id=REV_B, + source_file_name=FILE_B, + parse_track="chunk", + created_at=now, + updated_at=now, + ), + ] + ) + + def section( + section_id: str, + document_id: str, + job_result_id: str, + path: str, + title: str, + level: int, + sort_order: int, + parent: str | None, + summary: str = "", + ) -> DocumentSection: + return DocumentSection( + section_id=section_id, + user_id=USER_ID, + namespace=NAMESPACE, + document_id=document_id, + job_result_id=job_result_id, + parent_section_id=parent, + section_path=path, + section_title=title, + section_level=level, + summary=summary, + sort_order=sort_order, + created_at=now, + ) + + def chunk( + row_id: str, + chunk_id: str, + document_id: str, + job_result_id: str, + section_id: str, + content: str, + *, + chunk_type: str = "text", + connect_to: list[dict[str, str]] | None = None, + file_path: str | None = None, + ) -> DocumentChunk: + return DocumentChunk( + id=row_id, + chunk_id=chunk_id, + user_id=USER_ID, + namespace=NAMESPACE, + document_id=document_id, + job_result_id=job_result_id, + section_id=section_id, + chunk_type=chunk_type, + content=content, + term_search_text=content, + file_path=file_path, + source_chunk_path=file_path or chunk_id, + chunk_metadata={"connect_to": connect_to or [], "summary": content}, + sort_order=0, + created_at=now, + ) + + session.add_all( + [ + section("sec_root", DOC_A, REV_A, PATH_ROOT, "Root", 0, 0, None), + section("sec_overview", DOC_A, REV_A, PATH_OVERVIEW, "1 Overview", 1, 1, "sec_root"), + section( + "sec_findings", + DOC_A, + REV_A, + PATH_FINDINGS, + "1.1 Findings", + 2, + 2, + "sec_overview", + "hit-marker-findings " + "long summary text " * 12, + ), + section("sec_detail", DOC_A, REV_A, PATH_DETAIL, "Detail", 3, 3, "sec_findings"), + section("sec_treatment", DOC_A, REV_A, PATH_TREATMENT, "2 Treatment", 1, 4, "sec_root"), + section("sec_root_b", DOC_B, REV_B, f"{FILE_B} / Root", "Root", 0, 0, None), + section("sec_intro", DOC_B, REV_B, PATH_INTRO, "Intro", 1, 1, "sec_root_b"), + chunk("row_findings", "chunk_findings", DOC_A, REV_A, "sec_findings", "findings body"), + chunk("row_detail", "chunk_detail", DOC_A, REV_A, "sec_detail", "detail body"), + chunk("row_treatment", "chunk_treatment", DOC_A, REV_A, "sec_treatment", "treatment body"), + chunk("row_intro", "chunk_intro", DOC_B, REV_B, "sec_intro", "intro body"), + chunk( + "row_host_overview", + "chunk_host_overview", + DOC_A, + REV_A, + "sec_overview", + "overview hosts table", + connect_to=[{"target": "chunk_table_overview", "ref": "[tables/a.html]"}], + ), + chunk( + "row_host_treatment", + "chunk_host_treatment", + DOC_A, + REV_A, + "sec_treatment", + "treatment hosts table", + connect_to=[{"target": "chunk_table_treatment", "ref": "[tables/b.html]"}], + ), + chunk( + "row_table_overview", + "chunk_table_overview", + DOC_A, + REV_A, + "sec_root", + "overview table", + chunk_type="table", + file_path="tables/a.html", + ), + chunk( + "row_table_treatment", + "chunk_table_treatment", + DOC_A, + REV_A, + "sec_root", + "treatment table", + chunk_type="table", + file_path="tables/b.html", + ), + ] + ) + session.commit() + + +@pytest.fixture +def map_ctx() -> ToolContext: + engine = create_engine("sqlite:///:memory:") + + @event.listens_for(engine, "connect") + def _disable_fk(dbapi_connection, _connection_record) -> None: # noqa: ANN001 + dbapi_connection.execute("PRAGMA foreign_keys=OFF") + + JobResult.__table__.create(engine) + Document.__table__.create(engine) + DocumentSection.__table__.create(engine) + DocumentChunk.__table__.create(engine) + session = Session(engine) + _seed(session) + return ToolContext( + db=_AsyncSessionAdapter(session), # type: ignore[arg-type] + user_id=USER_ID, + namespace=NAMESPACE, + db_factory=_unused_db_factory, + query="findings", + ) + + +@pytest.mark.asyncio +async def test_outline_rejects_a_leaf_section(map_ctx: ToolContext) -> None: + result = await outline( + map_ctx, {"scope": [{"document_id": DOC_A, "section_path": PATH_TREATMENT}]} + ) + assert result.error is not None + assert "leaf section" in result.error + assert "corpus.read" in result.error + + +@pytest.mark.asyncio +async def test_unknown_scope_path_fails_the_call(map_ctx: ToolContext) -> None: + result = await outline( + map_ctx, {"scope": [{"document_id": DOC_A, "section_path": "no such path"}]} + ) + assert result.error is not None + assert "unknown section_path" in result.error + assert "no such path" in result.error + + +@pytest.mark.asyncio +async def test_outline_accepts_several_scope_targets(map_ctx: ToolContext) -> None: + result = await outline( + map_ctx, + {"scope": [{"document_id": DOC_A}, {"document_id": DOC_B}]}, + ) + assert result.error is None + _assert_shared_search_payload(result) + paths = {row["section_path"] for row in result.payload["rows"]} + assert PATH_OVERVIEW in paths + assert PATH_INTRO in paths + assert "| document_id=" in result.text + + +@pytest.mark.asyncio +async def test_node_filter_keeps_ancestors_and_full_descendant_subtree( + map_ctx: ToolContext, +) -> None: + result = await node_filter( + map_ctx, + { + "scope": [{"document_id": DOC_A}], + "predicates": [{"field": "summary", "terms": ["hit-marker-findings"]}], + }, + ) + assert result.error is None + assert "[Hit]" in result.text + assert PATH_FINDINGS in result.text + assert PATH_ROOT in result.text + assert PATH_OVERVIEW in result.text + assert PATH_DETAIL in result.text + assert PATH_TREATMENT not in result.text + _assert_shared_search_payload(result) + hits = [row for row in result.payload["rows"] if row["is_hit"]] + assert [row["section_path"] for row in hits] == [PATH_FINDINGS] + + +@pytest.mark.asyncio +async def test_outline_over_budget_drops_summaries_before_anything_else( + map_ctx: ToolContext, monkeypatch: pytest.MonkeyPatch +) -> None: + whole = await outline(map_ctx, {"scope": [{"document_id": DOC_A}]}) + assert whole.error is None + assert "summary: hit-marker-findings" in whole.text + assert "summaries omitted" not in whole.text + + monkeypatch.setattr( + "shared.services.retrieval.agent_tools.map_render.MAP_TOOL_CHAR_BUDGET", + len(whole.text) - 10, + ) + map_ctx.query = "" + result = await outline(map_ctx, {"scope": [{"document_id": DOC_A}]}) + assert result.error is None + assert "summary:" not in result.text + assert "summaries omitted" in result.text + assert PATH_TREATMENT in result.text + assert all(row["summary"] == "" for row in result.payload["rows"]) + + +@pytest.mark.asyncio +async def test_outline_too_large_without_query_points_to_grep_recall_assets( + map_ctx: ToolContext, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + "shared.services.retrieval.agent_tools.map_render.MAP_TOOL_CHAR_BUDGET", + 80, + ) + map_ctx.query = "" + result = await outline(map_ctx, {"scope": [{"document_id": DOC_A}]}) + assert result.error is not None + assert "no user query" in result.error + assert "corpus.grep" in result.error + assert "corpus.recall" in result.error + assert "corpus.assets" in result.error + + +@pytest.mark.asyncio +async def test_recall_passes_section_subtree_into_discovery( + map_ctx: ToolContext, monkeypatch: pytest.MonkeyPatch +) -> None: + captured: dict[str, Any] = {} + + async def _discovery(*_args: object, **kwargs: object) -> DiscoveryResult: + captured.update(kwargs) + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + + monkeypatch.setattr( + "shared.services.retrieval.agent_tools.tools.recall.map_unit_discovery", + _discovery, + ) + result = await recall( + map_ctx, + {"query": "findings", "scope": [{"document_id": DOC_A, "section_path": PATH_OVERVIEW}]}, + ) + assert result.error is None + assert captured["section_targets"] == [(DOC_A, PATH_OVERVIEW)] + + +@pytest.mark.asyncio +async def test_grep_subtree_scope_excludes_sibling_sections(map_ctx: ToolContext) -> None: + result = await grep( + map_ctx, + { + "pattern": "body", + "scope": [{"document_id": DOC_A, "section_path": PATH_OVERVIEW}], + }, + ) + assert result.error is None + _assert_shared_search_payload(result) + chunk_ids = [ref["chunk_id"] for ref in result.refs] + assert "chunk_findings" in chunk_ids + assert "chunk_detail" in chunk_ids + assert "chunk_treatment" not in chunk_ids + assert "| document_id=" in result.text + + +@pytest.mark.asyncio +async def test_assets_subtree_scope_keeps_only_connected_in_scope_assets( + map_ctx: ToolContext, +) -> None: + result = await assets( + map_ctx, + {"scope": [{"document_id": DOC_A, "section_path": PATH_OVERVIEW}], "type": "table"}, + ) + assert result.error is None + _assert_shared_search_payload(result) + chunk_ids = [row["chunk_id"] for row in result.payload["rows"]] + assert chunk_ids == ["chunk_table_overview"] + assert "chunk_table_treatment" not in result.text + + +def test_fold_at_map_tool_budget_hides_low_score_keeps_hit_chain() -> None: + hit = MapNode("findings", "overview", render="H" * 8000, depth=2, score=5.0) + ancestor = MapNode("overview", "root", render="A" * 100, depth=1, score=5.0) + root = MapNode("root", None, render="R" * 50, score=5.0) + low = MapNode("treatment", "root", render="L" * 15000, depth=1, score=0.0) + folded = fold_map_nodes( + [root, ancestor, hit, low], + char_budget=MAP_TOOL_CHAR_BUDGET, + keep_ids={"findings"}, + ) + assert not folded.overflow + assert {"findings", "overview", "root"} <= folded.visible_ids + assert "treatment" not in folded.visible_ids + assert folded.hidden_count == 1 + + +def test_fold_ties_hide_larger_then_deeper_subtree_first() -> None: + root = MapNode("root", None, render="R" * 10) + small = MapNode("small", "root", render="S" * 100, depth=1) + big = MapNode("big", "root", render="B" * 100, depth=1) + big_child = MapNode("big_child", "big", render="C" * 100, depth=2) + folded = fold_map_nodes( + [root, small, big, big_child], + char_budget=250, + keep_ids={"root"}, + ) + assert folded.visible_ids == {"root", "small"} diff --git a/packages/shared-python/shared/tests/test_agent_tools_read.py b/packages/shared-python/shared/tests/test_agent_tools_read.py index 673766bdc..a4a56ee45 100644 --- a/packages/shared-python/shared/tests/test_agent_tools_read.py +++ b/packages/shared-python/shared/tests/test_agent_tools_read.py @@ -10,6 +10,7 @@ import os from contextlib import asynccontextmanager from datetime import datetime +from typing import Any os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") @@ -25,7 +26,7 @@ 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 -from shared.services.retrieval.agent_tools.tools.read import read +from shared.services.retrieval.agent_tools.tools.read import _PICK_REMINDER, read USER_ID = "user_read" NAMESPACE = "default" @@ -231,18 +232,53 @@ def _read_kwargs(refs: list[dict[str, str]], *, mode: str = "self") -> dict: } +def _ref_status( + *, + document_id: str, + status: str, + section_path: str | None = None, + chunk_id: str | None = None, + reason: str | None = None, + chunk_ids: list[str] | None = None, +) -> dict[str, Any]: + return { + "document_id": document_id, + "section_path": section_path, + "chunk_id": chunk_id, + "status": status, + "reason": reason, + "chunk_ids": chunk_ids or [], + } + + +def _refs_text(entries: list[dict[str, Any]]) -> str: + lines = ["refs:"] + for entry in entries: + tag = "[ok]" if entry["status"] == "ok" else f"[failed: {entry['reason']}]" + lines.append( + f" {tag} document_id={entry['document_id']} " + f"section_path={entry['section_path']} chunk_id={entry['chunk_id']}" + ) + return "\n".join(lines) + + def _assert_result( ctx: ToolContext, result, *, - text: str, - errors: list[str], + body_text: str, + ref_status: list[dict[str, Any]], refs: list[dict[str, str]], executes: int, ) -> None: assert result.error is None - assert result.text == text - assert result.payload["errors"] == errors + expected_text = ( + f"{_refs_text(ref_status)}\n{_PICK_REMINDER}\n{body_text}\n{_PICK_REMINDER}" + if body_text + else f"{_refs_text(ref_status)}\n{_PICK_REMINDER}" + ) + assert result.text == expected_text + assert result.payload["refs"] == ref_status assert result.refs == refs assert [row["chunk_id"] for row in result.payload["chunks"]] == [ ref["chunk_id"] for ref in refs @@ -264,13 +300,26 @@ async def test_read_multi_exact_section_paths_preserve_ref_order(read_ctx: ToolC _assert_result( read_ctx, result, - text=( + body_text=( f"### {FILE_A} ({DOC_A}) / {PATH_TREATMENT} [text]\n" "treatment body\n" f"### {FILE_A} ({DOC_A}) / {PATH_OVERVIEW} [text]\n" "overview body" ), - errors=[], + ref_status=[ + _ref_status( + document_id=DOC_A, + status="ok", + section_path=PATH_TREATMENT, + chunk_ids=[CHUNK_TREATMENT], + ), + _ref_status( + document_id=DOC_A, + status="ok", + section_path=PATH_OVERVIEW, + chunk_ids=[CHUNK_OVERVIEW], + ), + ], refs=[ {"document_id": DOC_A, "chunk_id": CHUNK_TREATMENT}, {"document_id": DOC_A, "chunk_id": CHUNK_OVERVIEW}, @@ -293,13 +342,26 @@ async def test_read_suffix_unique_and_cross_document(read_ctx: ToolContext) -> N _assert_result( read_ctx, result, - text=( + body_text=( f"### {FILE_B} ({DOC_B}) / {PATH_UNIQUE} [text]\n" "unique body\n" f"### {FILE_A} ({DOC_A}) / {PATH_TREATMENT} [text]\n" "treatment body" ), - errors=[], + ref_status=[ + _ref_status( + document_id=DOC_B, + status="ok", + section_path="Unique Leaf", + chunk_ids=[CHUNK_UNIQUE], + ), + _ref_status( + document_id=DOC_A, + status="ok", + section_path=PATH_TREATMENT, + chunk_ids=[CHUNK_TREATMENT], + ), + ], refs=[ {"document_id": DOC_B, "chunk_id": CHUNK_UNIQUE}, {"document_id": DOC_A, "chunk_id": CHUNK_TREATMENT}, @@ -321,18 +383,29 @@ async def test_read_ambiguous_suffix_does_not_contaminate_other_refs( ] ), ) - assert "ambiguous section_path '1.1 Findings'" in result.payload["errors"][0] - assert PATH_FINDINGS in result.payload["errors"][0] - assert PATH_ANNEX_FINDINGS in result.payload["errors"][0] + failed_reason = result.payload["refs"][0]["reason"] + assert failed_reason is not None + assert "ambiguous section_path '1.1 Findings'" in failed_reason + assert PATH_FINDINGS in failed_reason + assert PATH_ANNEX_FINDINGS in failed_reason _assert_result( read_ctx, result, - text=( - f"errors: {result.payload['errors'][0]}\n" - f"### {FILE_B} ({DOC_B}) / {PATH_INTRO} [text]\n" - "intro body" - ), - errors=result.payload["errors"], + body_text=(f"### {FILE_B} ({DOC_B}) / {PATH_INTRO} [text]\n" "intro body"), + ref_status=[ + _ref_status( + document_id=DOC_A, + status="failed", + section_path="1.1 Findings", + reason=failed_reason, + ), + _ref_status( + document_id=DOC_B, + status="ok", + section_path=PATH_INTRO, + chunk_ids=[CHUNK_INTRO], + ), + ], refs=[{"document_id": DOC_B, "chunk_id": CHUNK_INTRO}], executes=7, ) @@ -355,15 +428,26 @@ async def test_read_unknown_path_and_unknown_document_keep_valid_ref( _assert_result( read_ctx, result, - text=( - f"errors: unknown section_path for {DOC_A}: no such path; " - "unknown document_id: doc_missing\n" - f"### {FILE_B} ({DOC_B}) / {PATH_INTRO} [text]\n" - "intro body" - ), - errors=[ - f"unknown section_path for {DOC_A}: no such path", - "unknown document_id: doc_missing", + body_text=(f"### {FILE_B} ({DOC_B}) / {PATH_INTRO} [text]\n" "intro body"), + ref_status=[ + _ref_status( + document_id=DOC_A, + status="failed", + section_path="no such path", + reason=f"unknown section_path for {DOC_A}: no such path", + ), + _ref_status( + document_id="doc_missing", + status="failed", + section_path=PATH_INTRO, + reason="unknown document_id: doc_missing", + ), + _ref_status( + document_id=DOC_B, + status="ok", + section_path=PATH_INTRO, + chunk_ids=[CHUNK_INTRO], + ), ], refs=[{"document_id": DOC_B, "chunk_id": CHUNK_INTRO}], executes=7, @@ -385,14 +469,32 @@ async def test_read_chunk_id_interleaved_with_section_path(read_ctx: ToolContext _assert_result( read_ctx, result, - text=( - f"errors: unknown chunk_id: missing_chunk in {DOC_A}\n" + body_text=( f"### {FILE_A} ({DOC_A}) / {PATH_FINDINGS} [text]\n" "findings body\n" f"### {FILE_B} ({DOC_B}) / {PATH_INTRO} [text]\n" "intro body" ), - errors=[f"unknown chunk_id: missing_chunk in {DOC_A}"], + ref_status=[ + _ref_status( + document_id=DOC_A, + status="ok", + chunk_id=CHUNK_FINDINGS, + chunk_ids=[CHUNK_FINDINGS], + ), + _ref_status( + document_id=DOC_B, + status="ok", + section_path=PATH_INTRO, + chunk_ids=[CHUNK_INTRO], + ), + _ref_status( + document_id=DOC_A, + status="failed", + chunk_id="missing_chunk", + reason=f"unknown chunk_id: missing_chunk in {DOC_A}", + ), + ], refs=[ {"document_id": DOC_A, "chunk_id": CHUNK_FINDINGS}, {"document_id": DOC_B, "chunk_id": CHUNK_INTRO}, @@ -415,13 +517,20 @@ async def test_read_descendants_does_not_include_sibling_suffix_or_old_revision( _assert_result( read_ctx, result, - text=( + body_text=( f"### {FILE_A} ({DOC_A}) / {PATH_FINDINGS} [text]\n" "findings body\n" f"### {FILE_A} ({DOC_A}) / {PATH_OVERVIEW} [text]\n" "overview body" ), - errors=[], + ref_status=[ + _ref_status( + document_id=DOC_A, + status="ok", + section_path=PATH_OVERVIEW, + chunk_ids=[CHUNK_FINDINGS, CHUNK_OVERVIEW], + ), + ], refs=[ {"document_id": DOC_A, "chunk_id": CHUNK_FINDINGS}, {"document_id": DOC_A, "chunk_id": CHUNK_OVERVIEW}, @@ -447,13 +556,26 @@ async def test_read_duplicate_section_refs_emit_twice(read_ctx: ToolContext) -> _assert_result( read_ctx, result, - text=( + body_text=( f"### {FILE_A} ({DOC_A}) / {PATH_TREATMENT} [text]\n" "treatment body\n" f"### {FILE_A} ({DOC_A}) / {PATH_TREATMENT} [text]\n" "treatment body" ), - errors=[], + ref_status=[ + _ref_status( + document_id=DOC_A, + status="ok", + section_path=PATH_TREATMENT, + chunk_ids=[CHUNK_TREATMENT], + ), + _ref_status( + document_id=DOC_A, + status="ok", + section_path=PATH_TREATMENT, + chunk_ids=[CHUNK_TREATMENT], + ), + ], refs=[ {"document_id": DOC_A, "chunk_id": CHUNK_TREATMENT}, {"document_id": DOC_A, "chunk_id": CHUNK_TREATMENT}, diff --git a/packages/shared-python/shared/tests/test_agent_tools_read_normalize.py b/packages/shared-python/shared/tests/test_agent_tools_read_normalize.py deleted file mode 100644 index 50f759e9c..000000000 --- a/packages/shared-python/shared/tests/test_agent_tools_read_normalize.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Pure tests for ``corpus.read``'s ``_normalize_read_refs`` argument shim. - -Covers the live-observed malformed shape (flat ``document_id`` + -``section_paths``/``chunk_ids`` instead of nested ``refs``) alongside the -canonical shape, without touching the DB-backed ``read`` tool function -itself. -""" - -from __future__ import annotations - -from shared.services.retrieval.agent_tools.tools.read import _normalize_read_refs - - -def test_normalize_read_refs_passes_through_canonical_refs() -> None: - args = {"refs": [{"document_id": "doc_1", "section_path": "A / B"}]} - assert _normalize_read_refs(args) == args["refs"] - - -def test_normalize_read_refs_ignores_flat_shape_when_refs_present() -> None: - args = { - "refs": [{"document_id": "doc_1", "chunk_id": "chunk_1"}], - "document_id": "doc_2", - "section_paths": ["should not be used"], - } - assert _normalize_read_refs(args) == args["refs"] - - -def test_normalize_read_refs_flat_document_id_with_section_paths_list() -> None: - # The exact malformed shape observed live (debug_agent_explore_episode.py, - # cursor_sdk harness, 2026-09-10): document_id hoisted to the top level - # alongside a plural section_paths array instead of nested refs. - args = { - "document_id": "doc_1cec16fef768", - "section_paths": [ - "基层心血管病综合管理实践指南2020 / 3 危险因素干预", - "基层心血管病综合管理实践指南2020 / 4 疾病干预", - ], - } - assert _normalize_read_refs(args) == [ - { - "document_id": "doc_1cec16fef768", - "section_path": "基层心血管病综合管理实践指南2020 / 3 危险因素干预", - }, - { - "document_id": "doc_1cec16fef768", - "section_path": "基层心血管病综合管理实践指南2020 / 4 疾病干预", - }, - ] - - -def test_normalize_read_refs_flat_document_id_with_singular_section_path() -> None: - args = {"document_id": "doc_1", "section_path": "A / B"} - assert _normalize_read_refs(args) == [{"document_id": "doc_1", "section_path": "A / B"}] - - -def test_normalize_read_refs_flat_document_id_with_chunk_ids_list() -> None: - args = {"document_id": "doc_1", "chunk_ids": ["c1", "c2"]} - assert _normalize_read_refs(args) == [ - {"document_id": "doc_1", "chunk_id": "c1"}, - {"document_id": "doc_1", "chunk_id": "c2"}, - ] - - -def test_normalize_read_refs_flat_document_id_with_singular_chunk_id() -> None: - args = {"document_id": "doc_1", "chunk_id": "c1"} - assert _normalize_read_refs(args) == [{"document_id": "doc_1", "chunk_id": "c1"}] - - -def test_normalize_read_refs_combines_all_flat_variants() -> None: - args = { - "document_id": "doc_1", - "section_path": "A", - "section_paths": ["B"], - "chunk_id": "c1", - "chunk_ids": ["c2"], - } - assert _normalize_read_refs(args) == [ - {"document_id": "doc_1", "section_path": "A"}, - {"document_id": "doc_1", "section_path": "B"}, - {"document_id": "doc_1", "chunk_id": "c1"}, - {"document_id": "doc_1", "chunk_id": "c2"}, - ] - - -def test_normalize_read_refs_no_document_id_returns_empty() -> None: - assert _normalize_read_refs({"section_paths": ["A"]}) == [] - - -def test_normalize_read_refs_empty_args_returns_empty() -> None: - assert _normalize_read_refs({}) == [] - - -def test_normalize_read_refs_blank_strings_are_skipped() -> None: - args = { - "document_id": "doc_1", - "section_path": " ", - "section_paths": ["", " ", "A"], - } - assert _normalize_read_refs(args) == [{"document_id": "doc_1", "section_path": "A"}] diff --git a/packages/shared-python/shared/tests/test_agent_tools_read_trace.py b/packages/shared-python/shared/tests/test_agent_tools_read_trace.py new file mode 100644 index 000000000..a7c97cf09 --- /dev/null +++ b/packages/shared-python/shared/tests/test_agent_tools_read_trace.py @@ -0,0 +1,165 @@ +"""read per-ref status is what the episode trace records.""" + +from __future__ import annotations + +import os +from contextlib import asynccontextmanager +from datetime import datetime + +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 sqlalchemy import create_engine, event +from sqlalchemy.orm import Session + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.agent_explore.bridge import build_decision_trace +from shared.services.retrieval.agent_explore.shared import tool_message_content +from shared.services.retrieval.agent_explore.types import AgentStep +from shared.services.retrieval.agent_tools.registry import ToolContext +from shared.services.retrieval.agent_tools.tools.read import read + +USER_ID = "user_trace" +NAMESPACE = "default" +DOC_A = "doc_a" +REV_A = "jr_a" +FILE_A = "guide.pdf" +PATH_INTRO = f"{FILE_A} / Intro" +CHUNK_INTRO = "chunk_intro" + + +class _AsyncSessionAdapter: + def __init__(self, session: Session) -> None: + self._session = session + + async def execute(self, statement): # noqa: ANN001 + return self._session.execute(statement) + + +@asynccontextmanager +async def _unused_db_factory(): + raise AssertionError("read-trace test should not open a second session") + yield # pragma: no cover + + +@pytest.fixture +def read_ctx() -> ToolContext: + engine = create_engine("sqlite:///:memory:") + + @event.listens_for(engine, "connect") + def _disable_fk(dbapi_connection, _connection_record) -> None: # noqa: ANN001 + dbapi_connection.execute("PRAGMA foreign_keys=OFF") + + now = datetime(2026, 1, 1) + JobResult.__table__.create(engine) + Document.__table__.create(engine) + DocumentSection.__table__.create(engine) + DocumentChunk.__table__.create(engine) + session = Session(engine) + session.add_all( + [ + JobResult( + id=REV_A, job_id="job_a", delivery_mode="inline", created_at=now, updated_at=now + ), + Document( + document_id=DOC_A, + user_id=USER_ID, + namespace=NAMESPACE, + status="active", + current_job_result_id=REV_A, + source_file_name=FILE_A, + parse_track="chunk", + created_at=now, + updated_at=now, + ), + DocumentSection( + section_id="sec_intro", + user_id=USER_ID, + namespace=NAMESPACE, + document_id=DOC_A, + job_result_id=REV_A, + section_path=PATH_INTRO, + section_level=1, + sort_order=0, + created_at=now, + ), + DocumentChunk( + id="row_intro", + chunk_id=CHUNK_INTRO, + user_id=USER_ID, + namespace=NAMESPACE, + document_id=DOC_A, + job_result_id=REV_A, + section_id="sec_intro", + chunk_type="text", + content="intro body", + source_chunk_path=CHUNK_INTRO, + created_at=now, + ), + ] + ) + session.commit() + return ToolContext( + db=_AsyncSessionAdapter(session), # type: ignore[arg-type] + user_id=USER_ID, + namespace=NAMESPACE, + db_factory=_unused_db_factory, + ) + + +@pytest.mark.asyncio +async def test_read_partial_failure_status_is_in_payload_and_trace( + read_ctx: ToolContext, +) -> None: + result = await read( + read_ctx, + { + "refs": [ + {"document_id": DOC_A, "section_path": PATH_INTRO}, + {"document_id": DOC_A, "chunk_id": "missing_chunk"}, + ], + "mode": "self", + "include_assets": False, + "resolve_same_as": False, + }, + ) + assert result.error is None + statuses = result.payload["refs"] + assert statuses[0]["status"] == "ok" + assert statuses[1]["status"] == "failed" + assert "unknown chunk_id" in statuses[1]["reason"] + assert result.refs == [{"document_id": DOC_A, "chunk_id": CHUNK_INTRO}] + + observation = tool_message_content( + result, tool_name="corpus.read", max_chars=12_000 + ) + assert "[ok]" in observation + assert "[failed: unknown chunk_id: missing_chunk in doc_a]" in observation + + trace = build_decision_trace( + [ + AgentStep( + step_index=0, + tool_name="corpus.read", + tool_args={"refs": [{"document_id": DOC_A, "section_path": PATH_INTRO}]}, + observation_text=observation, + error=result.error, + elapsed_ms=1, + tokens_used_delta=0, + tokens_used_total=0, + ref_status=result.payload["refs"], + ) + ] + ) + recorded = trace[0].observation["observation_text"] + assert "[ok]" in recorded + assert "[failed: unknown chunk_id: missing_chunk in doc_a]" in recorded + assert trace[0].observation["ref_status"] == result.payload["refs"] + assert trace[0].result["status"] == "ok" + assert trace[0].result["error"] is None diff --git a/packages/shared-python/shared/tests/test_map_unit_section_subtree.py b/packages/shared-python/shared/tests/test_map_unit_section_subtree.py new file mode 100644 index 000000000..f4dd2a750 --- /dev/null +++ b/packages/shared-python/shared/tests/test_map_unit_section_subtree.py @@ -0,0 +1,28 @@ +"""Strict section-subtree clause for map-unit discovery (recall scope).""" + +from __future__ import annotations + +from shared.services.retrieval.search.map_unit_discovery import ( + _build_section_subtree_clause, +) + + +def test_section_subtree_clause_empty_when_no_named_path() -> None: + assert _build_section_subtree_clause(None) == ("", {}) + assert _build_section_subtree_clause([]) == ("", {}) + assert _build_section_subtree_clause([("doc_a", None)]) == ("", {}) + + +def test_section_subtree_clause_is_exact_path_or_children() -> None: + clause, params = _build_section_subtree_clause( + [("doc_a", "guide.pdf / 1 Overview"), ("doc_b", None)] + ) + assert clause.startswith("AND (") + assert "ds.section_path = :_sec_path_0" in clause + assert "ds.section_path LIKE :_sec_pathlike_0" in clause + assert params["_sec_doc_0"] == "doc_a" + assert params["_sec_path_0"] == "guide.pdf / 1 Overview" + assert params["_sec_pathlike_0"] == "guide.pdf / 1 Overview / %" + assert "(dmu.document_id = :_sec_doc_1)" in clause + assert params["_sec_doc_1"] == "doc_b" + assert "%" not in params["_sec_path_0"] diff --git a/scripts/publication_benchmark/state_snapshot.py b/scripts/publication_benchmark/state_snapshot.py index bf36beef9..6c784a2fb 100644 --- a/scripts/publication_benchmark/state_snapshot.py +++ b/scripts/publication_benchmark/state_snapshot.py @@ -386,7 +386,7 @@ def capture_semantic_fingerprints( db, "SELECT u.id, u.document_id, u.job_result_id, u.unit_id, u.section_id, " "u.unit_kind, u.path_token_count, u.content_token_count, " - "u.term_search_text_lower, u.has_image, u.has_table, u.sort_order " + "u.has_image, u.has_table, u.sort_order " "FROM document_map_units u JOIN documents d ON d.document_id = u.document_id " "WHERE d.user_id = :user_id AND d.namespace = :namespace", parameters, @@ -411,7 +411,6 @@ def capture_semantic_fingerprints( "unit_id": _unit_id(row["unit_id"], section_paths), "path_token_count": row["path_token_count"], "content_token_count": row["content_token_count"], - "term_search_text_lower": row["term_search_text_lower"], "has_image": row["has_image"], "has_table": row["has_table"], } diff --git a/scripts/retire_legacy_token_lookup.py b/scripts/retire_legacy_token_lookup.py new file mode 100644 index 000000000..4d709311b --- /dev/null +++ b/scripts/retire_legacy_token_lookup.py @@ -0,0 +1,74 @@ +"""Inspect, explicitly retire, or restore the legacy token lookup index.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +from sqlalchemy import create_engine, text +from sqlalchemy.engine import Engine, URL, make_url +from sqlalchemy.pool import NullPool + +from shared.services.retrieval.maintenance.legacy_token_lookup import ( + LegacyTokenLookupMaintenance, +) + + +def main() -> int: + parser: argparse.ArgumentParser = argparse.ArgumentParser(description=__doc__) + actions: argparse._MutuallyExclusiveGroup = parser.add_mutually_exclusive_group() + actions.add_argument("--apply", action="store_true", help="Retire the legacy index") + actions.add_argument("--restore", action="store_true", help="Rebuild the legacy index") + parser.add_argument("--statement-timeout-seconds", type=int, default=1800) + arguments: argparse.Namespace = parser.parse_args() + database_url: str | None = os.environ.get("DATABASE_URL") + if not database_url: + parser.error("DATABASE_URL must be set in the environment") + if arguments.statement_timeout_seconds <= 0: + parser.error("--statement-timeout-seconds must be positive") + url: URL = make_url(database_url).set(drivername="postgresql+psycopg2") + ssl_variables: dict[str, str] = { + "DB_SSL_MODE": "sslmode", + "DB_SSL_CERT": "sslcert", + "DB_SSL_KEY": "sslkey", + "DB_SSL_ROOT_CERT": "sslrootcert", + } + connection_arguments: dict[str, str] = { + parameter: os.environ[variable] + for variable, parameter in ssl_variables.items() + if os.environ.get(variable) + } + engine: Engine = create_engine( + url, connect_args=connection_arguments, poolclass=NullPool, + isolation_level="AUTOCOMMIT", + ) + try: + with engine.connect() as connection: + connection.execute(text("SET lock_timeout = '5s'")) + connection.execute(text( + f"SET statement_timeout = '{arguments.statement_timeout_seconds}s'" + )) + maintenance: LegacyTokenLookupMaintenance = LegacyTokenLookupMaintenance( + connection + ) + if arguments.apply: + maintenance.apply_retirement() + elif arguments.restore: + maintenance.restore_index() + print(json.dumps(maintenance.inspect_indexes(), indent=2, sort_keys=True)) + except RuntimeError as error: + print(str(error), file=sys.stderr) + return 1 + except Exception as error: + # Connection exceptions may include credentials. Report only their type. + print(f"Index maintenance failed: {type(error).__name__}", file=sys.stderr) + return 1 + finally: + engine.dispose() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/uv.lock b/uv.lock index 93ceab7e8..873fcd412 100644 --- a/uv.lock +++ b/uv.lock @@ -1499,6 +1499,7 @@ dependencies = [ { name = "gevent" }, { name = "httpx" }, { name = "jieba" }, + { name = "jsonschema" }, { name = "kombu" }, { name = "loguru" }, { name = "openai" }, @@ -1546,6 +1547,7 @@ requires-dist = [ { name = "gevent", specifier = ">=24.11.1" }, { name = "httpx", specifier = "==0.28.1" }, { name = "jieba", specifier = "==0.42.1" }, + { name = "jsonschema", specifier = ">=4.20.0" }, { name = "kombu", specifier = "==5.5.4" }, { name = "loguru", specifier = "==0.7.3" }, { name = "openai", specifier = ">=1.0.0" },