cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

cognee/tests/test_feedback_enrichment.py

@@ -0,0 +1,174 @@
+"""
+End-to-end integration test for feedback enrichment feature.
+
+Tests the complete feedback enrichment pipeline:
+1. Add data and cognify
+2. Run search with save_interaction=True to create CogneeUserInteraction nodes
+3. Submit feedback to create CogneeUserFeedback nodes
+4. Run memify with feedback enrichment tasks to create FeedbackEnrichment nodes
+5. Verify all nodes and edges are properly created and linked in the graph
+"""
+
+import os
+import pathlib
+from collections import Counter
+
+import cognee
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.pipelines.tasks.task import Task
+from cognee.modules.search.types import SearchType
+from cognee.shared.data_models import KnowledgeGraph
+from cognee.shared.logging_utils import get_logger
+from cognee.tasks.feedback.create_enrichments import create_enrichments
+from cognee.tasks.feedback.extract_feedback_interactions import (
+    extract_feedback_interactions,
+)
+from cognee.tasks.feedback.generate_improved_answers import generate_improved_answers
+from cognee.tasks.feedback.link_enrichments_to_feedback import (
+    link_enrichments_to_feedback,
+)
+from cognee.tasks.graph import extract_graph_from_data
+from cognee.tasks.storage import add_data_points
+
+logger = get_logger()
+
+
+async def main():
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(
+                pathlib.Path(__file__).parent,
+                ".data_storage/test_feedback_enrichment",
+            )
+        ).resolve()
+    )
+    cognee_directory_path = str(
+        pathlib.Path(
+            os.path.join(
+                pathlib.Path(__file__).parent,
+                ".cognee_system/test_feedback_enrichment",
+            )
+        ).resolve()
+    )
+
+    cognee.config.data_root_directory(data_directory_path)
+    cognee.config.system_root_directory(cognee_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    dataset_name = "feedback_enrichment_test"
+
+    await cognee.add("Cognee turns documents into AI memory.", dataset_name)
+    await cognee.cognify([dataset_name])
+
+    question_text = "Say something."
+    result = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION,
+        query_text=question_text,
+        save_interaction=True,
+    )
+
+    assert len(result) > 0, "Search should return non-empty results"
+
+    feedback_text = "This answer was completely useless, my feedback is definitely negative."
+    await cognee.search(
+        query_type=SearchType.FEEDBACK,
+        query_text=feedback_text,
+        last_k=1,
+    )
+
+    graph_engine = await get_graph_engine()
+    nodes_before, edges_before = await graph_engine.get_graph_data()
+
+    interaction_nodes_before = [
+        (node_id, props)
+        for node_id, props in nodes_before
+        if props.get("type") == "CogneeUserInteraction"
+    ]
+    feedback_nodes_before = [
+        (node_id, props)
+        for node_id, props in nodes_before
+        if props.get("type") == "CogneeUserFeedback"
+    ]
+
+    edge_types_before = Counter(edge[2] for edge in edges_before)
+
+    assert len(interaction_nodes_before) >= 1, (
+        f"Expected at least 1 CogneeUserInteraction node, found {len(interaction_nodes_before)}"
+    )
+    assert len(feedback_nodes_before) >= 1, (
+        f"Expected at least 1 CogneeUserFeedback node, found {len(feedback_nodes_before)}"
+    )
+
+    for node_id, props in feedback_nodes_before:
+        sentiment = props.get("sentiment", "")
+        score = props.get("score", 0)
+        feedback_text = props.get("feedback", "")
+        logger.info(
+            "Feedback node created",
+            feedback=feedback_text,
+            sentiment=sentiment,
+            score=score,
+        )
+
+    assert edge_types_before.get("gives_feedback_to", 0) >= 1, (
+        f"Expected at least 1 'gives_feedback_to' edge, found {edge_types_before.get('gives_feedback_to', 0)}"
+    )
+
+    extraction_tasks = [Task(extract_feedback_interactions, last_n=5)]
+    enrichment_tasks = [
+        Task(generate_improved_answers, top_k=20),
+        Task(create_enrichments),
+        Task(
+            extract_graph_from_data,
+            graph_model=KnowledgeGraph,
+            task_config={"batch_size": 10},
+        ),
+        Task(add_data_points, task_config={"batch_size": 10}),
+        Task(link_enrichments_to_feedback),
+    ]
+
+    await cognee.memify(
+        extraction_tasks=extraction_tasks,
+        enrichment_tasks=enrichment_tasks,
+        data=[{}],
+        dataset="feedback_enrichment_test_memify",
+    )
+
+    nodes_after, edges_after = await graph_engine.get_graph_data()
+
+    enrichment_nodes = [
+        (node_id, props)
+        for node_id, props in nodes_after
+        if props.get("type") == "FeedbackEnrichment"
+    ]
+
+    assert len(enrichment_nodes) >= 1, (
+        f"Expected at least 1 FeedbackEnrichment node, found {len(enrichment_nodes)}"
+    )
+
+    for node_id, props in enrichment_nodes:
+        assert "text" in props, f"FeedbackEnrichment node {node_id} missing 'text' property"
+
+    enrichment_node_ids = {node_id for node_id, _ in enrichment_nodes}
+    edges_with_enrichments = [
+        edge
+        for edge in edges_after
+        if edge[0] in enrichment_node_ids or edge[1] in enrichment_node_ids
+    ]
+
+    assert len(edges_with_enrichments) >= 1, (
+        f"Expected enrichment nodes to have at least 1 edge, found {len(edges_with_enrichments)}"
+    )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    logger.info("All feedback enrichment tests passed successfully")
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

cognee/tests/test_kuzu.py

@@ -38,22 +38,35 @@ async def main():
 
         dataset_name = "cs_explanations"
 
-        explanation_file_path = os.path.join(
+        explanation_file_path_nlp = os.path.join(
             pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
         )
-        await cognee.add([explanation_file_path], dataset_name)
+        await cognee.add([explanation_file_path_nlp], dataset_name)
 
-        text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-        At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-        Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-        The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-        Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-        In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-        """
-        await cognee.add([text], dataset_name)
+        explanation_file_path_quantum = os.path.join(
+            pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+        )
+
+        from cognee.infrastructure.databases.graph import get_graph_engine
+
+        graph_engine = await get_graph_engine()
+
+        is_empty = await graph_engine.is_empty()
+
+        assert is_empty, "Kuzu graph database is not empty"
+
+        await cognee.add([explanation_file_path_quantum], dataset_name)
+
+        is_empty = await graph_engine.is_empty()
+
+        assert is_empty, "Kuzu graph database should be empty before cognify"
 
         await cognee.cognify([dataset_name])
 
+        is_empty = await graph_engine.is_empty()
+
+        assert not is_empty, "Kuzu graph database should not be empty"
+
         from cognee.infrastructure.databases.vector import get_vector_engine
 
         vector_engine = get_vector_engine()
@@ -61,7 +74,7 @@ async def main():
         random_node_name = random_node.payload["text"]
 
         search_results = await cognee.search(
-            query_type=SearchType.INSIGHTS, query_text=random_node_name
+            query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
         )
         assert len(search_results) != 0, "The search results list is empty."
         print("\n\nExtracted sentences are:\n")
@@ -117,11 +130,10 @@ async def main():
         assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
 
         await cognee.prune.prune_system(metadata=True)
-        from cognee.infrastructure.databases.graph import get_graph_engine
 
-        graph_engine = await get_graph_engine()
-        nodes, edges = await graph_engine.get_graph_data()
-        assert len(nodes) == 0 and len(edges) == 0, "Kuzu graph database is not empty"
+        is_empty = await graph_engine.is_empty()
+
+        assert is_empty, "Kuzu graph database is not empty"
 
     finally:
         # Ensure cleanup even if tests fail

cognee/tests/test_lancedb.py

@@ -131,20 +131,16 @@ async def main():
     dataset_name_1 = "natural_language"
     dataset_name_2 = "quantum"
 
-    explanation_file_path = os.path.join(
+    explanation_file_path_nlp = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
     )
-    await cognee.add([explanation_file_path], dataset_name_1)
+    await cognee.add([explanation_file_path_nlp], dataset_name_1)
 
-    text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-    At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-    Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-    The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-    Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-    In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-    """
+    explanation_file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
 
-    await cognee.add([text], dataset_name_2)
+    await cognee.add([explanation_file_path_quantum], dataset_name_2)
 
     await cognee.cognify([dataset_name_2, dataset_name_1])
 
@@ -157,7 +153,7 @@ async def main():
     random_node_name = random_node.payload["text"]
 
     search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
+        query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")

cognee/tests/test_library.py

@@ -6,6 +6,7 @@ from cognee.modules.search.operations import get_history
 from cognee.modules.users.methods import get_default_user
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.search.types import SearchType
+from cognee import update
 
 logger = get_logger()
 
@@ -42,7 +43,7 @@ async def main():
 
     await cognee.add([text], dataset_name)
 
-    await cognee.cognify([dataset_name])
+    cognify_run_info = await cognee.cognify([dataset_name])
 
     from cognee.infrastructure.databases.vector import get_vector_engine
 
@@ -51,7 +52,7 @@ async def main():
     random_node_name = random_node.payload["text"]
 
     search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
+        query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
@@ -77,6 +78,35 @@ async def main():
 
     assert len(history) == 6, "Search history is not correct."
 
+    # Test updating of documents
+    # Get Pipeline Run object
+    pipeline_run_obj = list(cognify_run_info.values())[0]
+    for data_item in pipeline_run_obj.data_ingestion_info:
+        # Update all documents in dataset to only contain Mark and Cindy information
+        await update(
+            dataset_id=pipeline_run_obj.dataset_id,
+            data_id=data_item["data_id"],
+            data="Mark met with Cindy at a cafe.",
+        )
+
+    search_results = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION, query_text="What information do you contain?"
+    )
+    assert "Mark" in search_results[0], (
+        "Failed to update document, no mention of Mark in search results"
+    )
+    assert "Cindy" in search_results[0], (
+        "Failed to update document, no mention of Cindy in search results"
+    )
+    assert "Artificial intelligence" not in search_results[0], (
+        "Failed to update document, Artificial intelligence still mentioned in search results"
+    )
+
+    # Test visualization
+    from cognee import visualize_graph
+
+    await visualize_graph()
+
     # Assert local data files are cleaned properly
     await cognee.prune.prune_data()
     data_root_directory = get_storage_config()["data_root_directory"]

cognee/tests/test_neo4j.py

@@ -32,23 +32,34 @@ async def main():
 
     dataset_name = "cs_explanations"
 
-    explanation_file_path = os.path.join(
+    explanation_file_path_nlp = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
     )
-    await cognee.add([explanation_file_path], dataset_name)
+    from cognee.infrastructure.databases.graph import get_graph_engine
+
+    graph_engine = await get_graph_engine()
+
+    is_empty = await graph_engine.is_empty()
 
-    text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-    At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-    Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-    The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-    Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-    In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-    """
+    assert is_empty, "Graph has to be empty"
+
+    await cognee.add([explanation_file_path_nlp], dataset_name)
+
+    explanation_file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
 
-    await cognee.add([text], dataset_name)
+    await cognee.add([explanation_file_path_quantum], dataset_name)
+    is_empty = await graph_engine.is_empty()
+
+    assert is_empty, "Graph has to be empty before cognify"
 
     await cognee.cognify([dataset_name])
 
+    is_empty = await graph_engine.is_empty()
+
+    assert not is_empty, "Graph shouldn't be empty"
+
     from cognee.infrastructure.databases.vector import get_vector_engine
 
     vector_engine = get_vector_engine()
@@ -56,7 +67,7 @@ async def main():
     random_node_name = random_node.payload["text"]
 
     search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
+        query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
@@ -121,11 +132,8 @@ async def main():
     assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
 
     await cognee.prune.prune_system(metadata=True)
-    from cognee.infrastructure.databases.graph import get_graph_engine
-
-    graph_engine = await get_graph_engine()
-    nodes, edges = await graph_engine.get_graph_data()
-    assert len(nodes) == 0 and len(edges) == 0, "Neo4j graph database is not empty"
+    is_empty = await graph_engine.is_empty()
+    assert is_empty, "Neo4j graph database is not empty"
 
 
 if __name__ == "__main__":

cognee/tests/test_neptune_analytics_vector.py

@@ -38,20 +38,16 @@ async def main():
 
     dataset_name = "cs_explanations"
 
-    explanation_file_path = os.path.join(
+    explanation_file_path_nlp = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
     )
-    await cognee.add([explanation_file_path], dataset_name)
+    await cognee.add([explanation_file_path_nlp], dataset_name)
 
-    text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-    At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-    Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-    The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-    Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-    In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-    """
+    explanation_file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
 
-    await cognee.add([text], dataset_name)
+    await cognee.add([explanation_file_path_quantum], dataset_name)
 
     await cognee.cognify([dataset_name])
 
@@ -60,7 +56,7 @@ async def main():
     random_node_name = random_node.payload["text"]
 
     search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
+        query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")

cognee/tests/test_permissions.py

@@ -34,25 +34,21 @@ async def main():
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
 
-    explanation_file_path = os.path.join(
+    explanation_file_path_nlp = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
     )
 
     # Add document for default user
-    await cognee.add([explanation_file_path], dataset_name="NLP")
+    await cognee.add([explanation_file_path_nlp], dataset_name="NLP")
     default_user = await get_default_user()
 
-    text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-    At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-    Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-    The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-    Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-    In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-    """
+    explanation_file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
 
     # Add document for test user
     test_user = await create_user("user@example.com", "example")
-    await cognee.add([text], dataset_name="QUANTUM", user=test_user)
+    await cognee.add([explanation_file_path_quantum], dataset_name="QUANTUM", user=test_user)
 
     nlp_cognify_result = await cognee.cognify(["NLP"], user=default_user)
     quantum_cognify_result = await cognee.cognify(["QUANTUM"], user=test_user)
@@ -101,7 +97,7 @@ async def main():
     add_error = False
     try:
         await cognee.add(
-            [explanation_file_path],
+            [explanation_file_path_nlp],
             dataset_name="QUANTUM",
             dataset_id=test_user_dataset_id,
             user=default_user,
@@ -143,7 +139,7 @@ async def main():
 
     # Add new data to test_users dataset from default_user
     await cognee.add(
-        [explanation_file_path],
+        [explanation_file_path_nlp],
         dataset_name="QUANTUM",
         dataset_id=test_user_dataset_id,
         user=default_user,
@@ -216,7 +212,7 @@ async def main():
     )
 
     # Try deleting data from test_user dataset with default_user after getting delete permission
-    # Get the dataset data to find the ID of the remaining data item (explanation_file_path)
+    # Get the dataset data to find the ID of the remaining data item (explanation_file_path_nlp)
     test_user_dataset_data = await get_dataset_data(test_user_dataset_id)
     explanation_file_data_id = test_user_dataset_data[0].id
 

cognee/tests/test_pgvector.py

@@ -141,10 +141,10 @@ async def main():
     dataset_name_1 = "natural_language"
     dataset_name_2 = "quantum"
 
-    explanation_file_path = os.path.join(
+    explanation_file_path_nlp = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
     )
-    await cognee.add([explanation_file_path], dataset_name_1)
+    await cognee.add([explanation_file_path_nlp], dataset_name_1)
 
     text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
     At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
@@ -167,7 +167,7 @@ async def main():
     random_node_name = random_node.payload["text"]
 
     search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
+        query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
@@ -202,7 +202,7 @@ async def main():
     history = await get_history(user.id)
     assert len(history) == 8, "Search history is not correct."
 
-    await test_local_file_deletion(text, explanation_file_path)
+    await test_local_file_deletion(text, explanation_file_path_nlp)
 
     await cognee.prune.prune_data()
     data_root_directory = get_storage_config()["data_root_directory"]

cognee/tests/test_remote_kuzu.py

@@ -42,19 +42,16 @@ async def main():
 
         dataset_name = "cs_explanations"
 
-        explanation_file_path = os.path.join(
+        explanation_file_path_nlp = os.path.join(
             pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
         )
-        await cognee.add([explanation_file_path], dataset_name)
+        await cognee.add([explanation_file_path_nlp], dataset_name)
 
-        text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-        At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-        Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-        The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-        Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-        In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-        """
-        await cognee.add([text], dataset_name)
+        explanation_file_path_quantum = os.path.join(
+            pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+        )
+
+        await cognee.add([explanation_file_path_quantum], dataset_name)
 
         await cognee.cognify([dataset_name])
 
@@ -65,7 +62,7 @@ async def main():
         random_node_name = random_node.payload["text"]
 
         search_results = await cognee.search(
-            query_type=SearchType.INSIGHTS, query_text=random_node_name
+            query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
         )
         assert len(search_results) != 0, "The search results list is empty."
         print("\n\nExtracted sentences are:\n")

cognee/tests/test_s3_file_storage.py

@@ -47,7 +47,7 @@ async def main():
     random_node_name = random_node.payload["text"]
 
     search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
+        query_type=SearchType.GRAPH_COMPLETION, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")

cognee/tests/test_search_db.py

@@ -1,3 +1,5 @@
+import pathlib
+import os
 import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
@@ -27,15 +29,11 @@ async def main():
     text_1 = """Germany is located in europe right next to the Netherlands"""
     await cognee.add(text_1, dataset_name)
 
-    text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-    At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-    Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-    The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-    Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-    In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-    """
+    explanation_file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
 
-    await cognee.add([text], dataset_name)
+    await cognee.add([explanation_file_path_quantum], dataset_name)
 
     await cognee.cognify([dataset_name])