diff --git a/application/prompt_client/prompt_client.py b/application/prompt_client/prompt_client.py index add7663f2..d9e151118 100644 --- a/application/prompt_client/prompt_client.py +++ b/application/prompt_client/prompt_client.py @@ -1193,7 +1193,7 @@ def get_id_of_most_similar_cre_paginated( max_similarity = -1 most_similar_index = 0 most_similar_id = "" - for page in range(starting_page, total_pages): + for page in range(starting_page, total_pages + 1): existing_cres, existing_cre_ids = self.__load_cre_embeddings(embeddings) similarities = cosine_similarity(embedding_array, existing_cres) @@ -1201,13 +1201,14 @@ def get_id_of_most_similar_cre_paginated( max_similarity = np.max(similarities) most_similar_index = np.argmax(similarities) most_similar_id = existing_cre_ids[most_similar_index] - ( - embeddings, - total_pages, - _, - ) = self.database.get_embeddings_by_doc_type_paginated( - cre_defs.Credoctypes.CRE.value, page=page - ) + if page < total_pages: + ( + embeddings, + total_pages, + _, + ) = self.database.get_embeddings_by_doc_type_paginated( + cre_defs.Credoctypes.CRE.value, page=page + 1 + ) if max_similarity < similarity_threshold: logger.info( @@ -1264,9 +1265,10 @@ def get_id_of_most_similar_node_paginated( most_similar_index = int(np.argmax(similarities)) most_similar_id = existing_standard_ids[most_similar_index] - embeddings, _, _ = self.database.get_embeddings_by_doc_type_paginated( - doc_type=cre_defs.Credoctypes.Standard.value, page=page - ) + if page < total_pages: + embeddings, _, _ = self.database.get_embeddings_by_doc_type_paginated( + doc_type=cre_defs.Credoctypes.Standard.value, page=page + 1 + ) if max_similarity < similarity_threshold: logger.info( f"there is no good standard candidate for this other standard section, returning nothing, max similarity was {max_similarity}" diff --git a/application/tests/prompt_client_pgvector_similarity_test.py b/application/tests/prompt_client_pgvector_similarity_test.py index 845f8000d..d6e6a32fe 100644 --- a/application/tests/prompt_client_pgvector_similarity_test.py +++ b/application/tests/prompt_client_pgvector_similarity_test.py @@ -1,6 +1,7 @@ """Chat/import similarity prefers pgvector when the DB reports it ready.""" import unittest +from typing import Any, Callable, Dict, List, Tuple from unittest.mock import MagicMock, patch from application.prompt_client import prompt_client as prompt_client_mod @@ -44,6 +45,111 @@ def test_cre_paginated_uses_pgvector_when_available(self) -> None: self.assertEqual(kwargs["id_column"], "cre_id") +def _paginated_side_effect( + pages: Dict[int, Dict[str, List[float]]], +) -> Callable[..., Tuple[Dict[str, List[float]], int, int]]: + """Build a ``get_embeddings_by_doc_type_paginated`` side_effect from a + ``{page_number: {id: embedding}}`` fixture. + + Mirrors the real method's contract (page is 1-indexed, ``total_pages`` is + the page count, missing/absent ``page`` defaults to 1 like the real + method's first call in ``get_id_of_most_similar_cre_paginated``). + """ + total_pages = len(pages) + + def _side_effect( + *args: Any, **kwargs: Any + ) -> Tuple[Dict[str, List[float]], int, int]: + page = kwargs.get("page", 1) + return pages.get(page, {}), total_pages, page + + return _side_effect + + +class PaginatedSimilarityFallbackTest(unittest.TestCase): + """Non-pgvector fallback path of the two ``_paginated`` similarity + lookups: every page must be visited exactly once, in order, including + the final one. Regression coverage for the page-alignment bug where the + fallback silently dropped trailing pages. + """ + + QUERY_EMBEDDING = [1.0, 0.0] + MATCHING_VECTOR = [1.0, 0.0] # cosine similarity 1.0 with the query + NOISE_VECTOR = [0.0, 1.0] # cosine similarity 0.0 with the query + + def _make_handler( + self, can_use_pgvector: bool, pages: Dict[int, Dict[str, List[float]]] + ) -> Tuple[prompt_client_mod.PromptHandler, MagicMock]: + database = MagicMock() + database.can_use_pgvector_similarity.return_value = can_use_pgvector + database.get_embeddings_by_doc_type_paginated.side_effect = ( + _paginated_side_effect(pages) + ) + handler = prompt_client_mod.PromptHandler.__new__( + prompt_client_mod.PromptHandler + ) + handler.database = database + return handler, database + + def test_node_paginated_finds_match_only_on_final_page(self) -> None: + # 3 pages; the only match is on the last one. The old implementation + # never processed it (it processed page 1 twice, page 2 once, and + # discarded page 3 after fetching it). + pages = { + 1: {"noise-1": self.NOISE_VECTOR}, + 2: {"noise-2": self.NOISE_VECTOR}, + 3: {"target-node": self.MATCHING_VECTOR}, + } + handler, database = self._make_handler(can_use_pgvector=False, pages=pages) + + result = handler.get_id_of_most_similar_node_paginated( + self.QUERY_EMBEDDING, similarity_threshold=0.5 + ) + + self.assertEqual(result, ("target-node", 1.0)) + database.find_most_similar_embedding_id.assert_not_called() + + def test_cre_paginated_finds_match_only_on_final_page(self) -> None: + # Same fixture shape for the CRE method, whose loop bound excluded + # the final page outright (range(starting_page, total_pages)). + pages = { + 1: {"noise-1": self.NOISE_VECTOR}, + 2: {"noise-2": self.NOISE_VECTOR}, + 3: {"target-cre": self.MATCHING_VECTOR}, + } + handler, database = self._make_handler(can_use_pgvector=False, pages=pages) + + result = handler.get_id_of_most_similar_cre_paginated( + self.QUERY_EMBEDDING, similarity_threshold=0.5 + ) + + self.assertEqual(result, ("target-cre", 1.0)) + database.find_most_similar_embedding_id.assert_not_called() + + def test_node_paginated_single_page(self) -> None: + pages = {1: {"only-node": self.MATCHING_VECTOR}} + handler, database = self._make_handler(can_use_pgvector=False, pages=pages) + + result = handler.get_id_of_most_similar_node_paginated( + self.QUERY_EMBEDDING, similarity_threshold=0.5 + ) + + self.assertEqual(result, ("only-node", 1.0)) + # Single page: no page beyond it should ever be requested. + self.assertEqual(database.get_embeddings_by_doc_type_paginated.call_count, 1) + + def test_cre_paginated_single_page(self) -> None: + pages = {1: {"only-cre": self.MATCHING_VECTOR}} + handler, database = self._make_handler(can_use_pgvector=False, pages=pages) + + result = handler.get_id_of_most_similar_cre_paginated( + self.QUERY_EMBEDDING, similarity_threshold=0.5 + ) + + self.assertEqual(result, ("only-cre", 1.0)) + self.assertEqual(database.get_embeddings_by_doc_type_paginated.call_count, 1) + + class FindMostSimilarEmbeddingIdResilienceTest(unittest.TestCase): def test_query_error_returns_no_match(self) -> None: from application.database.db import Node_collection