updated adapter tests to ensure documents are returned in similarity order (#172)

* updated tests to ensure documents are returned in similarity order

* updated  test

Author: epinzur

packages/graph-retriever/src/graph_retriever/content.py

@@ -18,9 +18,6 @@ class Content:
         The content.
     embedding :
         The embedding of the content.
-    score :
-        The similarity of the embedding to the query.
-        This is optional, and may not be set depending on the content.
     metadata :
         The metadata associated with the content.
     mime_type :
@@ -31,16 +28,14 @@ class Content:
     content: str
     embedding: list[float]
     metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
-
     mime_type: str = "text/plain"
-    score: float | None = None
 
     @staticmethod
     def new(
         id: str,
         content: str,
         embedding: list[float] | Callable[[str], list[float]],
-        score: float | None = None,
+        *,
         metadata: dict[str, Any] | None = None,
         mime_type: str = "text/plain",
     ) -> Content:
@@ -56,8 +51,6 @@ def new(
         embedding :
             The embedding, or a function to apply to the content to compute the
             embedding.
-        score :
-            The similarity of the embedding to the query.
         metadata :
             The metadata associated with the content.
         mime_type :
@@ -72,7 +65,6 @@ def new(
             id=id,
             content=content,
             embedding=embedding(content) if callable(embedding) else embedding,
-            score=score,
             metadata=metadata or {},
             mime_type=mime_type,
         )

packages/graph-retriever/src/graph_retriever/strategies/scored.py

@@ -66,3 +66,10 @@ def iteration(self, nodes: Iterable[Node], tracker: NodeTracker) -> None:
             node = highest.node
             node.extra_metadata["_score"] = highest.score
             limit -= tracker.select_and_traverse([node])
+
+    @override
+    def finalize_nodes(self, selected):
+        selected = sorted(
+            selected, key=lambda node: node.extra_metadata["_score"], reverse=True
+        )
+        return super().finalize_nodes(selected)

packages/graph-retriever/src/graph_retriever/testing/adapter_tests.py

@@ -8,6 +8,7 @@
 from graph_retriever import Content
 from graph_retriever.adapters import Adapter
 from graph_retriever.edges import Edge, IdEdge, MetadataEdge
+from graph_retriever.utils.math import cosine_similarity
 
 
 def assert_valid_result(content: Content):
@@ -40,6 +41,48 @@ def assert_ids_any_order(
     assert set(result_ids) == set(expected), "should contain exactly expected IDs"
 
 
+def cosine_similarity_scores(
+    adapter: Adapter, query_or_embedding: str | list[float], ids: list[str]
+) -> dict[str, float]:
+    """Return the cosine similarity scores for the given IDs and query embedding."""
+    if len(ids) == 0:
+        return {}
+
+    docs = adapter.get(ids)
+    found_ids = (d.id for d in docs)
+    assert set(ids) == set(found_ids), "can't find all IDs"
+
+    if isinstance(query_or_embedding, str):
+        query_embedding = adapter.search_with_embedding(query_or_embedding, k=0)[0]
+    else:
+        query_embedding = query_or_embedding
+
+    scores: list[float] = cosine_similarity(
+        [query_embedding],
+        [d.embedding for d in docs],
+    )[0]
+
+    return {doc.id: score for doc, score in zip(docs, scores)}
+
+
+def assert_ids_in_cosine_similarity_order(
+    results: Iterable[Content],
+    expected: list[str],
+    query_embedding: list[float],
+    adapter: Adapter,
+) -> None:
+    """Assert the results are valid and in cosine similarity order."""
+    assert_valid_results(results)
+    result_ids = [r.id for r in results]
+
+    similarity_scores = cosine_similarity_scores(adapter, query_embedding, expected)
+    expected = sorted(expected, key=lambda id: similarity_scores[id], reverse=True)
+
+    assert result_ids == expected, (
+        "should contain expected IDs in cosine similarity order"
+    )
+
+
 @dataclass(kw_only=True)
 class AdapterComplianceCase(abc.ABC):
     """
@@ -77,8 +120,40 @@ class GetCase(AdapterComplianceCase):
     GetCase(id="one", request=["boar"], expected=["boar"]),
     GetCase(
         id="many",
-        request=["boar", "chinchilla", "cobra"],
-        expected=["boar", "chinchilla", "cobra"],
+        request=[
+            "alligator",
+            "barracuda",
+            "chameleon",
+            "cobra",
+            "crocodile",
+            "dolphin",
+            "eel",
+            "fish",
+            "gecko",
+            "iguana",
+            "jellyfish",
+            "komodo dragon",
+            "lizard",
+            "manatee",
+            "narwhal",
+        ],
+        expected=[
+            "alligator",
+            "barracuda",
+            "chameleon",
+            "cobra",
+            "crocodile",
+            "dolphin",
+            "eel",
+            "fish",
+            "gecko",
+            "iguana",
+            "jellyfish",
+            "komodo dragon",
+            "lizard",
+            "manatee",
+            "narwhal",
+        ],
     ),
     GetCase(
         id="missing",
@@ -410,7 +485,7 @@ def expected(self, method: str, case: AdapterComplianceCase) -> list[str]:
 
         Generally, this should *not* change the expected results, unless the the
         adapter being tested uses wildly different distance metrics or a
-        different embedding. The `AnimalsEmbedding` is deterimistic and the
+        different embedding. The `AnimalsEmbedding` is deterministic and the
         results across vector stores should generally be deterministic and
         consistent.
 
@@ -469,7 +544,7 @@ def test_search_with_embedding(
             search_case.query, **search_case.kwargs
         )
         assert_is_embedding(embedding)
-        assert_ids_any_order(results, expected)
+        assert_ids_in_cosine_similarity_order(results, expected, embedding, adapter)
 
     async def test_asearch_with_embedding(
         self, adapter: Adapter, search_case: SearchCase
@@ -480,21 +555,21 @@ async def test_asearch_with_embedding(
             search_case.query, **search_case.kwargs
         )
         assert_is_embedding(embedding)
-        assert_ids_any_order(results, expected)
+        assert_ids_in_cosine_similarity_order(results, expected, embedding, adapter)
 
     def test_search(self, adapter: Adapter, search_case: SearchCase) -> None:
         """Run tests for `search`."""
         expected = self.expected("search", search_case)
         embedding, _ = adapter.search_with_embedding(search_case.query, k=0)
         results = adapter.search(embedding, **search_case.kwargs)
-        assert_ids_any_order(results, expected)
+        assert_ids_in_cosine_similarity_order(results, expected, embedding, adapter)
 
     async def test_asearch(self, adapter: Adapter, search_case: SearchCase) -> None:
         """Run tests for `asearch`."""
         expected = self.expected("asearch", search_case)
         embedding, _ = await adapter.asearch_with_embedding(search_case.query, k=0)
         results = await adapter.asearch(embedding, **search_case.kwargs)
-        assert_ids_any_order(results, expected)
+        assert_ids_in_cosine_similarity_order(results, expected, embedding, adapter)
 
     def test_adjacent(self, adapter: Adapter, adjacent_case: AdjacentCase) -> None:
         """Run tests for `adjacent."""
@@ -506,7 +581,7 @@ def test_adjacent(self, adapter: Adapter, adjacent_case: AdjacentCase) -> None:
             k=adjacent_case.k,
             filter=adjacent_case.filter,
         )
-        assert_ids_any_order(results, expected)
+        assert_ids_in_cosine_similarity_order(results, expected, embedding, adapter)
 
     async def test_aadjacent(
         self, adapter: Adapter, adjacent_case: AdjacentCase
@@ -520,4 +595,4 @@ async def test_aadjacent(
             k=adjacent_case.k,
             filter=adjacent_case.filter,
         )
-        assert_ids_any_order(results, expected)
+        assert_ids_in_cosine_similarity_order(results, expected, embedding, adapter)

packages/graph-retriever/src/graph_retriever/traversal.py

@@ -343,19 +343,18 @@ def _contents_to_new_nodes(
             c.id: c for c in contents if c.id not in self._discovered_node_ids
         }
 
-        # Compute scores (as needed).
-        if any(c.score is None for c in content_dict.values()):
-            scores = cosine_similarity(
-                [self.strategy._query_embedding],
-                [c.embedding for c in content_dict.values() if c.score is None],
-            )[0]
-        else:
-            scores = []
+        if len(content_dict) == 0:
+            return []
+
+        # Compute scores.
+        scores: list[float] = cosine_similarity(
+            [self.strategy._query_embedding],
+            [c.embedding for c in content_dict.values()],
+        )[0]
 
         # Create the nodes
-        scores_it = iter(scores)
         nodes = []
-        for content in content_dict.values():
+        for content, score in zip(content_dict.values(), scores):
             # Determine incoming/outgoing edges.
             edges = self.edge_function(content)
 
@@ -370,7 +369,6 @@ def _contents_to_new_nodes(
                     default=0,
                 )
 
-            score = content.score or next(scores_it)
             nodes.append(
                 Node(
                     id=content.id,

packages/graph-retriever/src/graph_retriever/utils/top_k.py

@@ -1,5 +1,4 @@
 from collections.abc import Iterable
-from typing import cast
 
 from graph_retriever.content import Content
 from graph_retriever.utils.math import cosine_similarity_top_k
@@ -12,7 +11,7 @@ def top_k(
     k: int,
 ) -> list[Content]:
     """
-    Select the top-k contents from the given contet.
+    Select the top-k contents from the given content.
 
     Parameters
     ----------
@@ -26,35 +25,32 @@ def top_k(
     Returns
     -------
     list[Content]
-        Top-K by similarity. All results will have their `score` set.
+        Top-K by similarity.
     """
     # TODO: Consider handling specially cases of already-sorted batches (merge).
     # TODO: Consider passing threshold here to limit results.
 
     # Use dicts to de-duplicate by ID. This ensures we choose the top K distinct
     # content (rather than K copies of the same content).
-    scored = {c.id: c for c in contents if c.score is not None}
-    unscored = {c.id: c for c in contents if c.score is None if c.id not in scored}
+    unscored = {c.id: c for c in contents}
 
-    if unscored:
-        top_unscored = _similarity_sort_top_k(
-            list(unscored.values()), embedding=embedding, k=k
-        )
-        scored.update(top_unscored)
+    top_scored = _similarity_sort_top_k(
+        list(unscored.values()), embedding=embedding, k=k
+    )
 
-    sorted = list(scored.values())
+    sorted = list(top_scored.values())
     sorted.sort(key=_score, reverse=True)
 
-    return sorted[:k]
+    return [c[0] for c in sorted]
 
 
-def _score(content: Content) -> float:
-    return cast(float, content.score)
+def _score(content_with_score: tuple[Content, float]) -> float:
+    return content_with_score[1]
 
 
 def _similarity_sort_top_k(
     contents: list[Content], *, embedding: list[float], k: int
-) -> dict[str, Content]:
+) -> dict[str, tuple[Content, float]]:
     # Flatten the content and use a dict to deduplicate.
     # We need to do this *before* selecting the top_k to ensure we don't
     # get duplicates (and fail to produce `k`).
@@ -65,6 +61,5 @@ def _similarity_sort_top_k(
     results = {}
     for (_x, y), score in zip(top_k, scores):
         c = contents[y]
-        c.score = score
-        results[c.id] = c
+        results[c.id] = (c, score)
     return results

packages/graph-retriever/tests/strategies/test_eager.py

@@ -6,6 +6,7 @@
 from graph_retriever.strategies import (
     Eager,
 )
+from graph_retriever.testing.adapter_tests import cosine_similarity_scores
 from graph_retriever.testing.embeddings import (
     ParserEmbeddings,
     angular_2d_embedding,
@@ -133,23 +134,17 @@ async def test_animals_habitat(animals: Adapter, sync_or_async: SyncOrAsync):
     ]
 
 
-async def test_animals_populates_metrics(animals: Adapter, sync_or_async: SyncOrAsync):
-    """Test that score and depth are populated."""
+async def test_animals_populates_metrics_and_order(
+    animals: Adapter, sync_or_async: SyncOrAsync
+):
+    """Test that score and depth are populated and results are returned in order."""
     results = await sync_or_async.traverse(
         store=animals,
         query=ANIMALS_QUERY,
         edges=[("habitat", "habitat")],
         strategy=Eager(select_k=100, start_k=2, max_depth=2),
     )()
 
-    expected_similarity_scores = {
-        "mongoose": 0.578682,
-        "bobcat": 0.02297939,
-        "cobra": 0.01365448699,
-        "deer": 0.1869947,
-        "elk": 0.02876833,
-        "fox": 0.533316,
-    }
     expected_depths = {
         "mongoose": 0,
         "bobcat": 1,
@@ -159,6 +154,10 @@ async def test_animals_populates_metrics(animals: Adapter, sync_or_async: SyncOr
         "fox": 0,
     }
 
+    expected_similarity_scores = cosine_similarity_scores(
+        animals, ANIMALS_QUERY, list(expected_depths.keys())
+    )
+
     for n in results:
         assert n.extra_metadata["_similarity_score"] == pytest.approx(
             expected_similarity_scores[n.id]
@@ -167,6 +166,15 @@ async def test_animals_populates_metrics(animals: Adapter, sync_or_async: SyncOr
             f"incorrect depth for {n.id}"
         )
 
+    expected_ids_in_order = sorted(
+        expected_similarity_scores.keys(),
+        key=lambda id: expected_similarity_scores[id],
+        reverse=True,
+    )
+    assert [n.id for n in results] == expected_ids_in_order, (
+        "incorrect order of results"
+    )
+
 
 async def test_animals_habitat_to_keywords(
     animals: Adapter, sync_or_async: SyncOrAsync