Remove dependency on precisely (#449)

Author: cbornet

libs/knowledge-graph/pyproject.toml

@@ -33,7 +33,6 @@ testcontainers = "~3.7.1"
 # https://github.com/psf/requests/issues/6707
 requests = "<=2.31.0"
 pytest = "^8.1.1"
-precisely = "^0.1.9"
 pytest-asyncio = "^0.23.6"
 pytest-dotenv = "^0.5.2"
 setuptools = "^70.0.0"

libs/knowledge-graph/tests/conftest.py

@@ -3,6 +3,7 @@
 
 import pytest
 from cassandra.cluster import Cluster, Session
+from dotenv import load_dotenv
 from langchain.graphs.graph_document import GraphDocument, Node, Relationship
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
@@ -11,6 +12,8 @@
 
 from ragstack_knowledge_graph.cassandra_graph_store import CassandraGraphStore
 
+load_dotenv()
+
 
 @pytest.fixture(scope="session")
 def db_keyspace() -> str:
@@ -65,7 +68,9 @@ def llm() -> BaseChatModel:
 
 
 class DataFixture:
-    def __init__(self, session: Session, keyspace: str, documents: List[GraphDocument]) -> None:
+    def __init__(
+        self, session: Session, keyspace: str, documents: List[GraphDocument]
+    ) -> None:
         self.session = session
         self.keyspace = "default_keyspace"
         self.uid = secrets.token_hex(8)
@@ -130,7 +135,9 @@ def marie_curie(db_session: Session, db_keyspace: str) -> Iterator[DataFixture]:
             Relationship(source=marie_curie, target=nobel_prize, type="WON"),
             Relationship(source=pierre_curie, target=nobel_prize, type="WON"),
             Relationship(source=marie_curie, target=pierre_curie, type="MARRIED_TO"),
-            Relationship(source=marie_curie, target=university_of_paris, type="WORKED_AT"),
+            Relationship(
+                source=marie_curie, target=university_of_paris, type="WORKED_AT"
+            ),
             Relationship(source=marie_curie, target=professor, type="HAS_PROFESSION"),
         ],
         source=Document(page_content="test_content"),

libs/knowledge-graph/tests/test_extraction.py

@@ -4,7 +4,6 @@
 from langchain_community.graphs.graph_document import Node, Relationship
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
-from precisely import assert_that, contains_exactly
 
 from ragstack_knowledge_graph.extraction import (
     KnowledgeSchema,
@@ -34,6 +33,7 @@ def extractor(llm: BaseChatModel) -> KnowledgeSchemaExtractor:
 Paris.
 """
 
+
 @pytest.mark.flaky(retries=10, delay=0)
 def test_extraction(extractor: KnowledgeSchemaExtractor):
     results = extractor.extract([Document(page_content=MARIE_CURIE_SOURCE)])
@@ -50,9 +50,8 @@ def test_extraction(extractor: KnowledgeSchemaExtractor):
     # putting things into standard title case, etc.
     university_of_paris = Node(id="University Of Paris", type="Institution")
 
-    assert_that(
-        results[0].nodes,
-        contains_exactly(
+    assert sorted(results[0].nodes, key=lambda x: x.id) == sorted(
+        [
             marie_curie,
             polish,
             french,
@@ -61,18 +60,24 @@ def test_extraction(extractor: KnowledgeSchemaExtractor):
             nobel_prize,
             pierre_curie,
             university_of_paris,
-        ),
+        ],
+        key=lambda x: x.id,
     )
-    assert_that(
-        results[0].relationships,
-        contains_exactly(
+
+    assert sorted(
+        results[0].relationships, key=lambda x: (x.source.id, x.target.id, x.type)
+    ) == sorted(
+        [
             Relationship(source=marie_curie, target=polish, type="HAS_NATIONALITY"),
             Relationship(source=marie_curie, target=french, type="HAS_NATIONALITY"),
             Relationship(source=marie_curie, target=physicist, type="HAS_OCCUPATION"),
             Relationship(source=marie_curie, target=chemist, type="HAS_OCCUPATION"),
             Relationship(source=marie_curie, target=nobel_prize, type="RECEIVED"),
             Relationship(source=pierre_curie, target=nobel_prize, type="RECEIVED"),
-            Relationship(source=marie_curie, target=university_of_paris, type="WORKED_AT"),
+            Relationship(
+                source=marie_curie, target=university_of_paris, type="WORKED_AT"
+            ),
             Relationship(source=marie_curie, target=pierre_curie, type="MARRIED_TO"),
-        ),
+        ],
+        key=lambda x: (x.source.id, x.target.id, x.type),
     )

libs/knowledge-graph/tests/test_knowledge_graph.py

@@ -1,13 +1,13 @@
 import secrets
 import pytest
-from precisely import assert_that, contains_exactly
 
 from cassandra.cluster import Session
 from ragstack_knowledge_graph.knowledge_graph import CassandraKnowledgeGraph
 from ragstack_knowledge_graph.traverse import Node, Relation
 
 from .conftest import DataFixture
 
+
 def test_no_embeddings(db_session: Session, db_keyspace: str) -> None:
     uid = secrets.token_hex(8)
     node_table = f"entities_{uid}"
@@ -22,6 +22,7 @@ def test_no_embeddings(db_session: Session, db_keyspace: str) -> None:
     )
     graph.insert([Node(name="a", type="b")])
 
+
 def test_traverse_marie_curie(marie_curie: DataFixture) -> None:
     (result_nodes, result_edges) = marie_curie.graph_store.graph.subgraph(
         start=Node("Marie Curie", "Person"),
@@ -40,45 +41,67 @@ def test_traverse_marie_curie(marie_curie: DataFixture) -> None:
         Node(name="Professor", type="Profession"),
     ]
     expected_edges = {
-        Relation(Node("Marie Curie", "Person"), Node("Polish", "Nationality"), "HAS_NATIONALITY"),
-        Relation(Node("Marie Curie", "Person"), Node("French", "Nationality"), "HAS_NATIONALITY"),
         Relation(
-            Node("Marie Curie", "Person"), Node("Physicist", "Profession"), "HAS_PROFESSION"
+            Node("Marie Curie", "Person"),
+            Node("Polish", "Nationality"),
+            "HAS_NATIONALITY",
+        ),
+        Relation(
+            Node("Marie Curie", "Person"),
+            Node("French", "Nationality"),
+            "HAS_NATIONALITY",
+        ),
+        Relation(
+            Node("Marie Curie", "Person"),
+            Node("Physicist", "Profession"),
+            "HAS_PROFESSION",
+        ),
+        Relation(
+            Node("Marie Curie", "Person"),
+            Node("Chemist", "Profession"),
+            "HAS_PROFESSION",
         ),
-        Relation(Node("Marie Curie", "Person"), Node("Chemist", "Profession"), "HAS_PROFESSION"),
         Relation(
-            Node("Marie Curie", "Person"), Node("Professor", "Profession"), "HAS_PROFESSION"
+            Node("Marie Curie", "Person"),
+            Node("Professor", "Profession"),
+            "HAS_PROFESSION",
         ),
         Relation(
             Node("Marie Curie", "Person"),
             Node("Radioactivity", "Scientific concept"),
             "RESEARCHED",
         ),
         Relation(Node("Marie Curie", "Person"), Node("Nobel Prize", "Award"), "WON"),
-        Relation(Node("Marie Curie", "Person"), Node("Pierre Curie", "Person"), "MARRIED_TO"),
+        Relation(
+            Node("Marie Curie", "Person"), Node("Pierre Curie", "Person"), "MARRIED_TO"
+        ),
         Relation(
             Node("Marie Curie", "Person"),
             Node("University of Paris", "Organization"),
             "WORKED_AT",
         ),
     }
-    assert_that(result_edges, contains_exactly(*expected_edges))
-    assert_that(result_nodes, contains_exactly(*expected_nodes))
+    assert sorted(result_edges) == sorted(expected_edges)
+    assert sorted(result_nodes) == sorted(expected_nodes)
 
 
 def test_fuzzy_search(marie_curie: DataFixture) -> None:
     if not marie_curie.has_embeddings:
-        pytest.skip("Fuzzy search requires embeddings. Run with openai environment variables")
-    result_nodes = marie_curie.graph_store.graph.query_nearest_nodes(["Marie", "Poland"])
+        pytest.skip(
+            "Fuzzy search requires embeddings. Run with openai environment variables"
+        )
+    result_nodes = marie_curie.graph_store.graph.query_nearest_nodes(
+        ["Marie", "Poland"]
+    )
     expected_nodes = [
         Node(name="Marie Curie", type="Person"),
         Node(name="Polish", type="Nationality", properties={"European": True}),
     ]
-    assert_that(result_nodes, contains_exactly(*expected_nodes))
+    assert sorted(result_nodes) == sorted(expected_nodes)
 
     result_nodes = marie_curie.graph_store.graph.query_nearest_nodes(["European"], k=2)
     expected_nodes = [
         Node(name="Polish", type="Nationality", properties={"European": True}),
         Node(name="French", type="Nationality", properties={"European": True}),
     ]
-    assert_that(result_nodes, contains_exactly(*expected_nodes))
+    assert sorted(result_nodes) == sorted(expected_nodes)

libs/knowledge-graph/tests/test_runnables.py

@@ -1,12 +1,9 @@
-from precisely import assert_that, contains_exactly
-
 from ragstack_knowledge_graph.runnables import extract_entities
 from ragstack_knowledge_graph.traverse import Node
 
 
 def test_extract_entities(llm):
     extractor = extract_entities(llm)
-    assert_that(
-        extractor.invoke({"question": "Who is Marie Curie?"}),
-        contains_exactly(Node("Marie Curie", "Person")),
-    )
+    assert extractor.invoke({"question": "Who is Marie Curie?"}) == [
+        Node("Marie Curie", "Person")
+    ]