diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index e5a4bad6..796d2ba8 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -6,7 +6,7 @@ from uuid import uuid4 from httpx import AsyncClient -from sqlalchemy import select +from sqlalchemy import Engine, event, select from shared.models.database.document import DocumentMapUnit from shared.services.retrieval.publication_content import ( @@ -88,6 +88,88 @@ async def test_classic_route_maps_winning_unit_to_one_chunk( } +async def test_classic_route_uses_token_hash_lookup_for_frequency_query( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-token-hash-{identifier}" + statements: list[str] = [] + + def capture_frequency_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if ( + "FROM document_map_unit_tokens AS tokens" in statement + and "tokens.frequency" in statement + ): + statements.append(statement) + + event.listen(Engine, "before_cursor_execute", capture_frequency_query) + try: + async with developer_api_client_factory() as api_client: + await _publish_document( + namespace=namespace, + source_file_name="token-hash.pdf", + chunks=[ + { + "chunk_id": f"token-hash-{identifier}", + "type": "text", + "content": "token hash lookup marker", + "path": "token-hash.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"token-hash-filler-a-{identifier}", + "type": "text", + "content": "unrelated filler a", + "path": "token-hash.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"token-hash-filler-b-{identifier}", + "type": "text", + "content": "unrelated filler b", + "path": "token-hash.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + { + "chunk_id": f"token-hash-filler-c-{identifier}", + "type": "text", + "content": "unrelated filler c", + "path": "token-hash.pdf/Root/Section/c", + "order": 4, + "metadata": {}, + }, + ], + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "token hash lookup", + "top_k": 1, + "use_agentic": False, + }, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_frequency_query) + + assert response.status_code == 200 + assert statements + assert "token_hash = ANY" in statements[-1] + assert "token = ANY" not in statements[-1] + + async def test_classic_route_image_filter_scores_only_units_with_images( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] 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 ab76977f..dc68d896 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 @@ -19,6 +19,7 @@ import time from collections.abc import Mapping from dataclasses import dataclass, field +from hashlib import sha256 from typing import Any from loguru import logger @@ -191,6 +192,9 @@ async def map_unit_discovery( query_tokens = tokenize_query_for_ranker(query) if not query_tokens: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + query_token_hashes = [ + sha256(token.encode("utf-8")).hexdigest() for token in query_tokens + ] revision_join, revision_clause, revision_params = _build_revision_scope( revision_pins @@ -245,13 +249,13 @@ async def map_unit_discovery( FROM document_map_unit_tokens AS tokens JOIN scoped_units ON scoped_units.map_unit_id = tokens.map_unit_id WHERE tokens.channel = ANY(:channels) - AND tokens.token = ANY(:tokens) + AND tokens.token_hash = ANY(:token_hashes) """ ), { **params, "channels": list(_MAP_SCORE_CHANNELS), - "tokens": query_tokens, + "token_hashes": query_token_hashes, }, ) frequencies: dict[tuple[str, str], dict[str, int]] = {}