cognee 0.5.0.dev0__py3-none-any.whl → 0.5.0.dev1__py3-none-any.whl

cognee/tests/integration/retrieval/test_triplet_retriever.py

@@ -0,0 +1,84 @@
+import os
+import pytest
+import pathlib
+import pytest_asyncio
+import cognee
+
+from cognee.low_level import setup
+from cognee.tasks.storage import add_data_points
+from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
+from cognee.modules.engine.models import Triplet
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment_with_triplets():
+    """Set up a clean test environment with triplets."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(base_dir / ".cognee_system/test_triplet_retriever_context_simple")
+    data_directory_path = str(base_dir / ".data_storage/test_triplet_retriever_context_simple")
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    await setup()
+
+    triplet1 = Triplet(
+        from_node_id="node1",
+        to_node_id="node2",
+        text="Alice knows Bob",
+    )
+    triplet2 = Triplet(
+        from_node_id="node2",
+        to_node_id="node3",
+        text="Bob works at Tech Corp",
+    )
+
+    triplets = [triplet1, triplet2]
+    await add_data_points(triplets)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment_empty():
+    """Set up a clean test environment without triplets."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(
+        base_dir / ".cognee_system/test_triplet_retriever_context_empty_collection"
+    )
+    data_directory_path = str(
+        base_dir / ".data_storage/test_triplet_retriever_context_empty_collection"
+    )
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_simple(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever can retrieve triplet context."""
+    retriever = TripletRetriever(top_k=5)
+
+    context = await retriever.get_context("Alice")
+
+    assert "Alice knows Bob" in context, "Failed to get Alice triplet"

cognee/tests/integration/tasks/test_add_data_points.py

@@ -0,0 +1,139 @@
+import pathlib
+import pytest
+import pytest_asyncio
+
+import cognee
+from cognee.low_level import setup
+from cognee.infrastructure.engine import DataPoint
+from cognee.tasks.storage.add_data_points import add_data_points
+from cognee.tasks.storage.exceptions import InvalidDataPointsInAddDataPointsError
+from cognee.infrastructure.databases.graph import get_graph_engine
+
+
+class Person(DataPoint):
+    name: str
+    age: int
+    metadata: dict = {"index_fields": ["name"]}
+
+
+class Company(DataPoint):
+    name: str
+    industry: str
+    metadata: dict = {"index_fields": ["name", "industry"]}
+
+
+@pytest_asyncio.fixture
+async def clean_test_environment():
+    """Set up a clean test environment for add_data_points tests."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(base_dir / ".cognee_system/test_add_data_points_integration")
+    data_directory_path = str(base_dir / ".data_storage/test_add_data_points_integration")
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    await setup()
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_add_data_points_comprehensive(clean_test_environment):
+    """Comprehensive integration test for add_data_points functionality."""
+
+    person1 = Person(name="Alice", age=30)
+    person2 = Person(name="Bob", age=25)
+    result = await add_data_points([person1, person2])
+
+    assert result == [person1, person2]
+    assert len(result) == 2
+
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) >= 2
+
+    result_empty = await add_data_points([])
+    assert result_empty == []
+
+    person3 = Person(name="Charlie", age=35)
+    person4 = Person(name="Diana", age=32)
+    custom_edge = (str(person3.id), str(person4.id), "knows", {"edge_text": "friends with"})
+
+    result_custom = await add_data_points([person3, person4], custom_edges=[custom_edge])
+    assert len(result_custom) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(edges) == 1
+    assert len(nodes) == 4
+
+    class Employee(DataPoint):
+        name: str
+        works_at: Company
+        metadata: dict = {"index_fields": ["name"]}
+
+    company = Company(name="TechCorp", industry="Technology")
+    employee = Employee(name="Eve", works_at=company)
+
+    result_rel = await add_data_points([employee])
+    assert len(result_rel) == 1
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 6
+    assert len(edges) == 2
+
+    person5 = Person(name="Frank", age=40)
+    person6 = Person(name="Grace", age=38)
+    triplet_edge = (str(person5.id), str(person6.id), "married_to", {"edge_text": "is married to"})
+
+    result_triplet = await add_data_points(
+        [person5, person6], custom_edges=[triplet_edge], embed_triplets=True
+    )
+    assert len(result_triplet) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 8
+    assert len(edges) == 3
+
+    batch1 = [Person(name="Leo", age=25), Person(name="Mia", age=30)]
+    batch2 = [Person(name="Noah", age=35), Person(name="Olivia", age=40)]
+
+    result_batch1 = await add_data_points(batch1)
+    result_batch2 = await add_data_points(batch2)
+
+    assert len(result_batch1) == 2
+    assert len(result_batch2) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 12
+    assert len(edges) == 3
+
+    person7 = Person(name="Paul", age=33)
+    person8 = Person(name="Quinn", age=31)
+    edge1 = (str(person7.id), str(person8.id), "colleague_of", {"edge_text": "works with"})
+    edge2 = (str(person8.id), str(person7.id), "colleague_of", {"edge_text": "works with"})
+
+    result_bi = await add_data_points([person7, person8], custom_edges=[edge1, edge2])
+    assert len(result_bi) == 2
+
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 14
+    assert len(edges) == 5
+
+    person_invalid = Person(name="Invalid", age=50)
+    with pytest.raises(InvalidDataPointsInAddDataPointsError, match="must be a list"):
+        await add_data_points(person_invalid)
+
+    with pytest.raises(InvalidDataPointsInAddDataPointsError, match="must be a DataPoint"):
+        await add_data_points(["not", "datapoints"])
+
+    final_nodes, final_edges = await graph_engine.get_graph_data()
+    assert len(final_nodes) == 14
+    assert len(final_edges) == 5

cognee/tests/integration/tasks/test_get_triplet_datapoints.py

@@ -0,0 +1,69 @@
+import os
+import pathlib
+import pytest
+import pytest_asyncio
+from unittest.mock import AsyncMock, patch
+
+import cognee
+from cognee.tasks.memify.get_triplet_datapoints import get_triplet_datapoints
+from cognee.modules.engine.models import Triplet
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment():
+    """Set up a clean test environment with a simple graph."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    data_directory_path = str(base_dir / ".data_storage/test_get_triplet_datapoints_integration")
+    cognee_directory_path = str(base_dir / ".cognee_system/test_get_triplet_datapoints_integration")
+
+    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 = "test_triplets"
+
+    text = "Volkswagen is a german car manufacturer from Wolfsburg. They produce different models such as Golf, Polo and Touareg."
+    await cognee.add(text, dataset_name)
+    await cognee.cognify([dataset_name])
+
+    yield dataset_name
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+
+@pytest.mark.asyncio
+async def test_get_triplet_datapoints_integration(setup_test_environment):
+    """Integration test: verify get_triplet_datapoints works with real graph data."""
+
+    from cognee.infrastructure.databases.graph import get_graph_engine
+
+    graph_engine = await get_graph_engine()
+
+    if not hasattr(graph_engine, "get_triplets_batch"):
+        pytest.skip("Graph engine does not support get_triplets_batch")
+
+    triplets = []
+    with patch(
+        "cognee.tasks.memify.get_triplet_datapoints.index_data_points", new_callable=AsyncMock
+    ):
+        async for triplet in get_triplet_datapoints([{}], triplets_batch_size=10):
+            triplets.append(triplet)
+
+    nodes, edges = await graph_engine.get_graph_data()
+
+    if len(edges) > 0 and len(triplets) == 0:
+        test_triplets = await graph_engine.get_triplets_batch(offset=0, limit=10)
+        if len(test_triplets) == 0:
+            pytest.fail(
+                f"Edges exist in graph ({len(edges)} edges) but get_triplets_batch found none. "
+                f"This indicates the query pattern may not match the graph structure."
+            )
+
+    for triplet in triplets:
+        assert isinstance(triplet, Triplet), "Each item should be a Triplet instance"
+        assert triplet.from_node_id, "Triplet should have from_node_id"
+        assert triplet.to_node_id, "Triplet should have to_node_id"
+        assert triplet.text, "Triplet should have embeddable text"

cognee/tests/test_cognee_server_start.py

@@ -25,8 +25,6 @@ class TestCogneeServerStart(unittest.TestCase):
                 "--port",
                 "8000",
             ],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
             preexec_fn=os.setsid,
         )
         # Give the server some time to start
@@ -150,8 +148,8 @@ class TestCogneeServerStart(unittest.TestCase):
             headers=headers,
             files=[("ontology_file", ("test.owl", ontology_content, "application/xml"))],
             data={
-                "ontology_key": json.dumps([ontology_key]),
-                "description": json.dumps(["Test ontology"]),
+                "ontology_key": ontology_key,
+                "description": "Test ontology",
             },
         )
         self.assertEqual(ontology_response.status_code, 200)

cognee/tests/test_conversation_history.py

@@ -8,10 +8,10 @@ Tests all retrievers that save conversation history to Redis cache:
 4. GRAPH_COMPLETION_CONTEXT_EXTENSION
 5. GRAPH_SUMMARY_COMPLETION
 6. TEMPORAL
+7. TRIPLET_COMPLETION
 """
 
 import os
-import shutil
 import cognee
 import pathlib
 
@@ -63,6 +63,10 @@ async def main():
 
     user = await get_default_user()
 
+    from cognee.memify_pipelines.create_triplet_embeddings import create_triplet_embeddings
+
+    await create_triplet_embeddings(user=user, dataset=dataset_name)
+
     cache_engine = get_cache_engine()
     assert cache_engine is not None, "Cache engine should be available for testing"
 
@@ -216,6 +220,24 @@ async def main():
     ]
     assert len(our_qa_temporal) == 1, "Should find Temporal question in history"
 
+    session_id_triplet = "test_session_triplet"
+
+    result_triplet = await cognee.search(
+        query_type=SearchType.TRIPLET_COMPLETION,
+        query_text="What companies are mentioned?",
+        session_id=session_id_triplet,
+    )
+
+    assert isinstance(result_triplet, list) and len(result_triplet) > 0, (
+        f"TRIPLET_COMPLETION should return non-empty list, got: {result_triplet!r}"
+    )
+
+    history_triplet = await cache_engine.get_latest_qa(str(user.id), session_id_triplet, last_n=10)
+    our_qa_triplet = [
+        h for h in history_triplet if h["question"] == "What companies are mentioned?"
+    ]
+    assert len(our_qa_triplet) == 1, "Should find Triplet question in history"
+
     from cognee.modules.retrieval.utils.session_cache import (
         get_conversation_history,
     )

cognee/tests/test_dataset_database_handler.py

@@ -0,0 +1,137 @@
+import asyncio
+import os
+
+# Set custom dataset database handler environment variable
+os.environ["VECTOR_DATASET_DATABASE_HANDLER"] = "custom_lancedb_handler"
+os.environ["GRAPH_DATASET_DATABASE_HANDLER"] = "custom_kuzu_handler"
+
+import cognee
+from cognee.modules.users.methods import get_default_user
+from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
+from cognee.shared.logging_utils import setup_logging, ERROR
+from cognee.api.v1.search import SearchType
+
+
+class LanceDBTestDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    @classmethod
+    async def create_dataset(cls, dataset_id, user):
+        import pathlib
+
+        cognee_directory_path = str(
+            pathlib.Path(
+                os.path.join(
+                    pathlib.Path(__file__).parent, ".cognee_system/test_dataset_database_handler"
+                )
+            ).resolve()
+        )
+        databases_directory_path = os.path.join(cognee_directory_path, "databases", str(user.id))
+        os.makedirs(databases_directory_path, exist_ok=True)
+
+        vector_db_name = "test.lance.db"
+
+        return {
+            "vector_dataset_database_handler": "custom_lancedb_handler",
+            "vector_database_name": vector_db_name,
+            "vector_database_url": os.path.join(databases_directory_path, vector_db_name),
+            "vector_database_provider": "lancedb",
+        }
+
+
+class KuzuTestDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    @classmethod
+    async def create_dataset(cls, dataset_id, user):
+        databases_directory_path = os.path.join("databases", str(user.id))
+        os.makedirs(databases_directory_path, exist_ok=True)
+
+        graph_db_name = "test.kuzu"
+        return {
+            "graph_dataset_database_handler": "custom_kuzu_handler",
+            "graph_database_name": graph_db_name,
+            "graph_database_url": os.path.join(databases_directory_path, graph_db_name),
+            "graph_database_provider": "kuzu",
+        }
+
+
+async def main():
+    import pathlib
+
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(
+                pathlib.Path(__file__).parent, ".data_storage/test_dataset_database_handler"
+            )
+        ).resolve()
+    )
+    cognee.config.data_root_directory(data_directory_path)
+    cognee_directory_path = str(
+        pathlib.Path(
+            os.path.join(
+                pathlib.Path(__file__).parent, ".cognee_system/test_dataset_database_handler"
+            )
+        ).resolve()
+    )
+    cognee.config.system_root_directory(cognee_directory_path)
+
+    # Add custom dataset database handler
+    from cognee.infrastructure.databases.dataset_database_handler.use_dataset_database_handler import (
+        use_dataset_database_handler,
+    )
+
+    use_dataset_database_handler(
+        "custom_lancedb_handler", LanceDBTestDatasetDatabaseHandler, "lancedb"
+    )
+    use_dataset_database_handler("custom_kuzu_handler", KuzuTestDatasetDatabaseHandler, "kuzu")
+
+    # Create a clean slate for cognee -- reset data and system state
+    print("Resetting cognee data...")
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    print("Data reset complete.\n")
+
+    # cognee knowledge graph will be created based on this text
+    text = """
+    Natural language processing (NLP) is an interdisciplinary
+    subfield of computer science and information retrieval.
+    """
+
+    print("Adding text to cognee:")
+    print(text.strip())
+
+    # Add the text, and make it available for cognify
+    await cognee.add(text)
+    print("Text added successfully.\n")
+
+    # Use LLMs and cognee to create knowledge graph
+    await cognee.cognify()
+    print("Cognify process complete.\n")
+
+    query_text = "Tell me about NLP"
+    print(f"Searching cognee for insights with query: '{query_text}'")
+    # Query cognee for insights on the added text
+    search_results = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION, query_text=query_text
+    )
+
+    print("Search results:")
+    # Display results
+    for result_text in search_results:
+        print(result_text)
+
+    default_user = await get_default_user()
+    # Assert that the custom database files were created based on the custom dataset database handlers
+    assert os.path.exists(
+        os.path.join(cognee_directory_path, "databases", str(default_user.id), "test.kuzu")
+    ), "Graph database file not found."
+    assert os.path.exists(
+        os.path.join(cognee_directory_path, "databases", str(default_user.id), "test.lance.db")
+    ), "Vector database file not found."
+
+
+if __name__ == "__main__":
+    logger = setup_logging(log_level=ERROR)
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        loop.run_until_complete(main())
+    finally:
+        loop.run_until_complete(loop.shutdown_asyncgens())

cognee/tests/test_dataset_delete.py

@@ -0,0 +1,76 @@
+import os
+import asyncio
+import pathlib
+from uuid import UUID
+
+import cognee
+from cognee.shared.logging_utils import setup_logging, ERROR
+from cognee.modules.data.methods.delete_dataset import delete_dataset
+from cognee.modules.data.methods.get_dataset import get_dataset
+from cognee.modules.users.methods import get_default_user
+
+
+async def main():
+    # Set data and system directory paths
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_dataset_delete")
+        ).resolve()
+    )
+    cognee.config.data_root_directory(data_directory_path)
+    cognee_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_dataset_delete")
+        ).resolve()
+    )
+    cognee.config.system_root_directory(cognee_directory_path)
+
+    # Create a clean slate for cognee -- reset data and system state
+    print("Resetting cognee data...")
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    print("Data reset complete.\n")
+
+    # cognee knowledge graph will be created based on this text
+    text = """
+    Natural language processing (NLP) is an interdisciplinary
+    subfield of computer science and information retrieval.
+    """
+
+    # Add the text, and make it available for cognify
+    await cognee.add(text, "nlp_dataset")
+    await cognee.add("Quantum computing is the study of quantum computers.", "quantum_dataset")
+
+    # Use LLMs and cognee to create knowledge graph
+    ret_val = await cognee.cognify()
+    user = await get_default_user()
+
+    for val in ret_val:
+        dataset_id = str(val)
+        vector_db_path = os.path.join(
+            cognee_directory_path, "databases", str(user.id), dataset_id + ".lance.db"
+        )
+        graph_db_path = os.path.join(
+            cognee_directory_path, "databases", str(user.id), dataset_id + ".pkl"
+        )
+
+        # Check if databases are properly created and exist before deletion
+        assert os.path.exists(graph_db_path), "Graph database file not found."
+        assert os.path.exists(vector_db_path), "Vector database file not found."
+
+        dataset = await get_dataset(user_id=user.id, dataset_id=UUID(dataset_id))
+        await delete_dataset(dataset)
+
+        # Confirm databases have been deleted
+        assert not os.path.exists(graph_db_path), "Graph database file found."
+        assert not os.path.exists(vector_db_path), "Vector database file found."
+
+
+if __name__ == "__main__":
+    logger = setup_logging(log_level=ERROR)
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        loop.run_until_complete(main())
+    finally:
+        loop.run_until_complete(loop.shutdown_asyncgens())