move documents to tests for simple test scenarios

packages/langchain-graph-retriever/src/langchain_graph_retriever/retrievers/__init__.py

@@ -2,4 +2,4 @@
 
 __all__ = [
     "GraphTraversalRetriever",
-]
+]

packages/langchain-graph-retriever/src/langchain_graph_retriever/retrievers/strategy/mmr.py

@@ -3,10 +3,11 @@
 from typing import Iterable
 
 import numpy as np
-from langchain_graph_retriever.utils.math import cosine_similarity
 from numpy.typing import NDArray
 from pydantic import Field
 
+from langchain_graph_retriever.utils.math import cosine_similarity
+
 from ..node import Node
 from .base import TraversalStrategy
 

packages/langchain-graph-retriever/tests/integration_tests/retrievers/conftest.py

@@ -1,13 +1,5 @@
 from tests.integration_tests.invoker import invoker
 from tests.integration_tests.retrievers.animal_docs import animal_docs, animal_store
-from tests.integration_tests.retrievers.earth_docs import (
-    earth_docs,
-    earth_store,
-)
-from tests.integration_tests.retrievers.parser_docs import (
-    graph_vector_store_docs,
-    parser_store,
-)
 
 # Imports for definitions.
 from tests.integration_tests.stores import (
@@ -22,11 +14,7 @@
     store_factory,
     store_param,
     enabled_stores,
-    earth_docs,
-    earth_store,
     animal_docs,
     animal_store,
-    graph_vector_store_docs,
-    parser_store,
     invoker,
 )

packages/langchain-graph-retriever/tests/integration_tests/retrievers/test_graph_traversal_eager.py

@@ -1,16 +1,18 @@
+import pytest
 from langchain_core.documents import Document
 from langchain_graph_retriever.retrievers.graph_traversal_retriever import (
     GraphTraversalRetriever,
 )
 from langchain_graph_retriever.retrievers.strategy.eager import (
     Eager,
 )
+from tests.embeddings.simple_embeddings import EarthEmbeddings, ParserEmbeddings
 from tests.integration_tests.assertions import assert_document_format, sorted_doc_ids
 from tests.integration_tests.retrievers.animal_docs import (
     ANIMALS_DEPTH_0_EXPECTED,
     ANIMALS_QUERY,
 )
-from tests.integration_tests.stores import StoreAdapter
+from tests.integration_tests.stores import StoreAdapter, StoreFactory
 
 
 async def test_animals_bidir_collection_eager(animal_store: StoreAdapter, invoker):
@@ -117,9 +119,72 @@ async def test_animals_item_to_collection(animal_store: StoreAdapter, invoker):
     assert sorted_doc_ids(docs) == ["bear", "bobcat", "caribou", "fox", "mongoose"]
 
 
-async def test_parser(parser_store: StoreAdapter, invoker):
+async def test_parser(
+    request: pytest.FixtureRequest, store_factory: StoreFactory, invoker
+):
+    """
+    This is a test of set of Documents to pre-populate,
+    a graph vector store with entries placed in a certain way.
+
+    Space of the entries (under Euclidean similarity):
+
+                      A0    (*)
+        ....        AL   AR       <....
+        :              |              :
+        :              |  ^           :
+        v              |  .           v
+                       |   :
+       TR              |   :          BL
+    T0   --------------x--------------   B0
+       TL              |   :          BR
+                       |   :
+                       |  .
+                       | .
+                       |
+                    FL   FR
+                      F0
+
+    the query point is meant to be at (*).
+    the A are bidirectionally with B
+    the A are outgoing to T
+    the A are incoming from F
+    The links are like: L with L, 0 with 0 and R with R.
+    """
+
+    docs_a = [
+        Document(id="AL", page_content="[-1, 9]", metadata={"label": "AL"}),
+        Document(id="A0", page_content="[0, 10]", metadata={"label": "A0"}),
+        Document(id="AR", page_content="[1, 9]", metadata={"label": "AR"}),
+    ]
+    docs_b = [
+        Document(id="BL", page_content="[9, 1]", metadata={"label": "BL"}),
+        Document(id="B0", page_content="[10, 0]", metadata={"label": "B0"}),
+        Document(id="BR", page_content="[9, -1]", metadata={"label": "BR"}),
+    ]
+    docs_f = [
+        Document(id="FL", page_content="[1, -9]", metadata={"label": "FL"}),
+        Document(id="F0", page_content="[0, -10]", metadata={"label": "F0"}),
+        Document(id="FR", page_content="[-1, -9]", metadata={"label": "FR"}),
+    ]
+    docs_t = [
+        Document(id="TL", page_content="[-9, -1]", metadata={"label": "TL"}),
+        Document(id="T0", page_content="[-10, 0]", metadata={"label": "T0"}),
+        Document(id="TR", page_content="[-9, 1]", metadata={"label": "TR"}),
+    ]
+    for doc_a, suffix in zip(docs_a, ["l", "0", "r"]):
+        doc_a.metadata["tag"] = f"ab_{suffix}"
+        doc_a.metadata["out"] = f"at_{suffix}"
+        doc_a.metadata["in"] = f"af_{suffix}"
+    for doc_b, suffix in zip(docs_b, ["l", "0", "r"]):
+        doc_b.metadata["tag"] = f"ab_{suffix}"
+    for doc_t, suffix in zip(docs_t, ["l", "0", "r"]):
+        doc_t.metadata["in"] = f"at_{suffix}"
+    for doc_f, suffix in zip(docs_f, ["l", "0", "r"]):
+        doc_f.metadata["out"] = f"af_{suffix}"
+    documents = docs_a + docs_b + docs_f + docs_t
+
     retriever = GraphTraversalRetriever(
-        store=parser_store,
+        store=store_factory.create(request, ParserEmbeddings(dimension=2), documents),
         edges=[("out", "in"), "tag"],
         strategy=Eager(k=10, start_k=2, max_depth=2),
     )
@@ -139,9 +204,31 @@ async def test_parser(parser_store: StoreAdapter, invoker):
     assert_document_format(docs[0])
 
 
-async def test_earth(earth_store: StoreAdapter, invoker):
+async def test_earth(
+    request: pytest.FixtureRequest, store_factory: StoreFactory, invoker
+):
+    greetings = Document(
+        id="greetings",
+        page_content="Typical Greetings",
+        metadata={
+            "incoming": "parent",
+        },
+    )
+
+    doc1 = Document(
+        id="doc1",
+        page_content="Hello World",
+        metadata={"outgoing": "parent", "keywords": ["greeting", "world"]},
+    )
+
+    doc2 = Document(
+        id="doc2",
+        page_content="Hello Earth",
+        metadata={"outgoing": "parent", "keywords": ["greeting", "earth"]},
+    )
+
     retriever = GraphTraversalRetriever(
-        store=earth_store,
+        store=store_factory.create(request, EarthEmbeddings(), [greetings, doc1, doc2]),
         edges=[("outgoing", "incoming"), "keywords"],
         strategy=Eager(k=10, start_k=2, max_depth=0),
     )