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

cognee/tests/integration/tasks/test_add_data_points.py

@@ -1,139 +0,0 @@
-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

@@ -1,69 +0,0 @@
-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_dataset_database_handler.py

@@ -1,137 +0,0 @@
-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

@@ -1,76 +0,0 @@
-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())

cognee/tests/test_edge_centered_payload.py

@@ -1,170 +0,0 @@
-"""
-End-to-end integration test for edge-centered payload and triplet embeddings.
-
-"""
-
-import os
-import pathlib
-import cognee
-from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.modules.search.types import SearchType
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
-from cognee.modules.ontology.ontology_config import Config
-
-logger = get_logger()
-
-text_data = """
-Apple is a technology company that produces the iPhone, iPad, and Mac computers.
-The company is known for its innovative products and ecosystem integration.
-
-Microsoft develops the Windows operating system and Office productivity suite.
-They are also major players in cloud computing with Azure.
-
-Google created the Android operating system and provides search engine services.
-The company is a leader in artificial intelligence and machine learning.
-"""
-
-ontology_content = """<?xml version="1.0"?>
-<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
-         xmlns:owl="http://www.w3.org/2002/07/owl#"
-         xmlns:rdfs="http://www.w3.org/2000/01/rdf-schema#"
-         xmlns="http://example.org/tech#"
-         xml:base="http://example.org/tech">
-
-    <owl:Ontology rdf:about="http://example.org/tech"/>
-
-    <!-- Classes -->
-    <owl:Class rdf:ID="Company"/>
-    <owl:Class rdf:ID="TechnologyCompany"/>
-    <owl:Class rdf:ID="Product"/>
-    <owl:Class rdf:ID="Software"/>
-    <owl:Class rdf:ID="Hardware"/>
-    <owl:Class rdf:ID="Service"/>
-
-    <rdf:Description rdf:about="#TechnologyCompany">
-        <rdfs:subClassOf rdf:resource="#Company"/>
-        <rdfs:comment>A company operating in the technology sector.</rdfs:comment>
-    </rdf:Description>
-
-    <rdf:Description rdf:about="#Software">
-        <rdfs:subClassOf rdf:resource="#Product"/>
-        <rdfs:comment>Software products and applications.</rdfs:comment>
-    </rdf:Description>
-
-    <rdf:Description rdf:about="#Hardware">
-        <rdfs:subClassOf rdf:resource="#Product"/>
-        <rdfs:comment>Physical hardware products.</rdfs:comment>
-    </rdf:Description>
-
-    <!-- Individuals -->
-    <TechnologyCompany rdf:ID="apple">
-        <rdfs:label>Apple</rdfs:label>
-    </TechnologyCompany>
-
-    <TechnologyCompany rdf:ID="microsoft">
-        <rdfs:label>Microsoft</rdfs:label>
-    </TechnologyCompany>
-
-    <TechnologyCompany rdf:ID="google">
-        <rdfs:label>Google</rdfs:label>
-    </TechnologyCompany>
-
-    <Hardware rdf:ID="iphone">
-        <rdfs:label>iPhone</rdfs:label>
-    </Hardware>
-
-    <Software rdf:ID="windows">
-        <rdfs:label>Windows</rdfs:label>
-    </Software>
-
-    <Software rdf:ID="android">
-        <rdfs:label>Android</rdfs:label>
-    </Software>
-
-</rdf:RDF>"""
-
-
-async def main():
-    data_directory_path = str(
-        pathlib.Path(
-            os.path.join(
-                pathlib.Path(__file__).parent,
-                ".data_storage/test_edge_centered_payload",
-            )
-        ).resolve()
-    )
-    cognee_directory_path = str(
-        pathlib.Path(
-            os.path.join(
-                pathlib.Path(__file__).parent,
-                ".cognee_system/test_edge_centered_payload",
-            )
-        ).resolve()
-    )
-
-    cognee.config.data_root_directory(data_directory_path)
-    cognee.config.system_root_directory(cognee_directory_path)
-
-    dataset_name = "tech_companies"
-
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-
-    await cognee.add(data=text_data, dataset_name=dataset_name)
-
-    import tempfile
-
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".owl", delete=False) as f:
-        f.write(ontology_content)
-        ontology_file_path = f.name
-
-    try:
-        logger.info(f"Loading ontology from: {ontology_file_path}")
-        config: Config = {
-            "ontology_config": {
-                "ontology_resolver": RDFLibOntologyResolver(ontology_file=ontology_file_path)
-            }
-        }
-
-        await cognee.cognify(datasets=[dataset_name], config=config)
-        graph_engine = await get_graph_engine()
-        nodes_phase2, edges_phase2 = await graph_engine.get_graph_data()
-
-        vector_engine = get_vector_engine()
-        triplets_phase2 = await vector_engine.search(
-            query_text="technology", limit=None, collection_name="Triplet_text"
-        )
-
-        assert len(triplets_phase2) == len(edges_phase2), (
-            f"Triplet embeddings and number of edges do not match. Vector db contains {len(triplets_phase2)} edge triplets while graph db contains {len(edges_phase2)} edges."
-        )
-
-        search_results_phase2 = await cognee.search(
-            query_type=SearchType.TRIPLET_COMPLETION,
-            query_text="What products does Apple make?",
-        )
-
-        assert search_results_phase2 is not None, (
-            "Search should return results for triplet embeddings in simple ontology use case."
-        )
-
-    finally:
-        if os.path.exists(ontology_file_path):
-            os.unlink(ontology_file_path)
-
-
-if __name__ == "__main__":
-    import asyncio
-    from cognee.shared.logging_utils import setup_logging
-
-    setup_logging()
-
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    try:
-        loop.run_until_complete(main())
-    finally:
-        loop.run_until_complete(loop.shutdown_asyncgens())
-        loop.close()