cognee 0.3.4.dev3__py3-none-any.whl → 0.3.5__py3-none-any.whl

cognee/tests/test_lancedb.py

@@ -0,0 +1,211 @@
+import os
+import pathlib
+
+import cognee
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.files.storage import get_storage_config
+from cognee.modules.data.models import Data
+from cognee.modules.users.methods import get_default_user
+from cognee.modules.search.types import SearchType
+from cognee.modules.search.operations import get_history
+
+logger = get_logger()
+
+
+async def test_local_file_deletion(data_text, file_location):
+    from sqlalchemy import select
+    import hashlib
+    from cognee.infrastructure.databases.relational import get_relational_engine
+
+    engine = get_relational_engine()
+
+    async with engine.get_async_session() as session:
+        # Get hash of data contents
+        encoded_text = data_text.encode("utf-8")
+        data_hash = hashlib.md5(encoded_text).hexdigest()
+        # Get data entry from database based on hash contents
+        data = (await session.scalars(select(Data).where(Data.content_hash == data_hash))).one()
+        assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
+            f"Data location doesn't exist: {data.raw_data_location}"
+        )
+        # Test deletion of data along with local files created by cognee
+        await engine.delete_data_entity(data.id)
+        assert not os.path.exists(data.raw_data_location.replace("file://", "")), (
+            f"Data location still exists after deletion: {data.raw_data_location}"
+        )
+
+    async with engine.get_async_session() as session:
+        # Get data entry from database based on file path
+        data = (
+            await session.scalars(select(Data).where(Data.raw_data_location == file_location))
+        ).one()
+        assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
+            f"Data location doesn't exist: {data.raw_data_location}"
+        )
+        # Test local files not created by cognee won't get deleted
+        await engine.delete_data_entity(data.id)
+        assert os.path.exists(data.raw_data_location.replace("file://", "")), (
+            f"Data location doesn't exists: {data.raw_data_location}"
+        )
+
+
+async def test_getting_of_documents(dataset_name_1):
+    # Test getting of documents for search per dataset
+    from cognee.modules.users.permissions.methods import get_document_ids_for_user
+
+    user = await get_default_user()
+    document_ids = await get_document_ids_for_user(user.id, [dataset_name_1])
+    assert len(document_ids) == 1, (
+        f"Number of expected documents doesn't match {len(document_ids)} != 1"
+    )
+
+    # Test getting of documents for search when no dataset is provided
+    user = await get_default_user()
+    document_ids = await get_document_ids_for_user(user.id)
+    assert len(document_ids) == 2, (
+        f"Number of expected documents doesn't match {len(document_ids)} != 2"
+    )
+
+
+async def test_vector_engine_search_none_limit():
+    file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
+
+    file_path_nlp = os.path.join(
+        pathlib.Path(__file__).parent,
+        "test_data/Natural_language_processing.txt",
+    )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add(file_path_quantum)
+
+    await cognee.add(file_path_nlp)
+
+    await cognee.cognify()
+
+    query_text = "Tell me about Quantum computers"
+
+    from cognee.infrastructure.databases.vector import get_vector_engine
+
+    vector_engine = get_vector_engine()
+
+    collection_name = "Entity_name"
+
+    query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
+
+    result = await vector_engine.search(
+        collection_name=collection_name, query_vector=query_vector, limit=None
+    )
+
+    # Check that we did not accidentally use any default value for limit
+    # in vector search along the way (like 5, 10, or 15)
+    assert len(result) > 15
+
+
+async def main():
+    cognee.config.set_vector_db_config(
+        {
+            "vector_db_provider": "lancedb",
+        }
+    )
+
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_lancedb")
+        ).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_lancedb")
+        ).resolve()
+    )
+    cognee.config.system_root_directory(cognee_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    dataset_name_1 = "natural_language"
+    dataset_name_2 = "quantum"
+
+    explanation_file_path = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
+    )
+    await cognee.add([explanation_file_path], 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.
+    """
+
+    await cognee.add([text], dataset_name_2)
+
+    await cognee.cognify([dataset_name_2, dataset_name_1])
+
+    from cognee.infrastructure.databases.vector import get_vector_engine
+
+    await test_getting_of_documents(dataset_name_1)
+
+    vector_engine = get_vector_engine()
+    random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0]
+    random_node_name = random_node.payload["text"]
+
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
+    assert len(search_results) != 0, "The search results list is empty."
+    print("\n\nExtracted sentences are:\n")
+    for result in search_results:
+        print(f"{result}\n")
+
+    search_results = await cognee.search(
+        query_type=SearchType.CHUNKS, query_text=random_node_name, datasets=[dataset_name_2]
+    )
+    assert len(search_results) != 0, "The search results list is empty."
+    print("\n\nExtracted chunks are:\n")
+    for result in search_results:
+        print(f"{result}\n")
+
+    graph_completion = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION,
+        query_text=random_node_name,
+        datasets=[dataset_name_2],
+    )
+    assert len(graph_completion) != 0, "Completion result is empty."
+    print("Completion result is:")
+    print(graph_completion)
+
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
+    assert len(search_results) != 0, "Query related summaries don't exist."
+    print("\n\nExtracted summaries are:\n")
+    for result in search_results:
+        print(f"{result}\n")
+
+    user = await get_default_user()
+    history = await get_history(user.id)
+    assert len(history) == 8, "Search history is not correct."
+
+    await cognee.prune.prune_data()
+    data_root_directory = get_storage_config()["data_root_directory"]
+    assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
+
+    await cognee.prune.prune_system(metadata=True)
+    connection = await vector_engine.get_connection()
+    tables_in_database = await connection.table_names()
+    assert len(tables_in_database) == 0, "LanceDB database is not empty"
+
+    await test_vector_engine_search_none_limit()
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

cognee/tests/test_pgvector.py

@@ -68,6 +68,44 @@ async def test_getting_of_documents(dataset_name_1):
     )
 
 
+async def test_vector_engine_search_none_limit():
+    file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
+
+    file_path_nlp = os.path.join(
+        pathlib.Path(__file__).parent,
+        "test_data/Natural_language_processing.txt",
+    )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add(file_path_quantum)
+
+    await cognee.add(file_path_nlp)
+
+    await cognee.cognify()
+
+    query_text = "Tell me about Quantum computers"
+
+    from cognee.infrastructure.databases.vector import get_vector_engine
+
+    vector_engine = get_vector_engine()
+
+    collection_name = "Entity_name"
+
+    query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
+
+    result = await vector_engine.search(
+        collection_name=collection_name, query_vector=query_vector, limit=None
+    )
+
+    # Check that we did not accidentally use any default value for limit
+    # in vector search along the way (like 5, 10, or 15)
+    assert len(result) > 15
+
+
 async def main():
     cognee.config.set_vector_db_config(
         {"vector_db_url": "", "vector_db_key": "", "vector_db_provider": "pgvector"}
@@ -174,6 +212,8 @@ async def main():
     tables_in_database = await vector_engine.get_table_names()
     assert len(tables_in_database) == 0, "PostgreSQL database is not empty"
 
+    await test_vector_engine_search_none_limit()
+
 
 if __name__ == "__main__":
     import asyncio

cognee/tests/test_relational_db_migration.py

@@ -197,6 +197,80 @@ async def relational_db_migration():
     print(f"All checks passed for {graph_db_provider} provider with '{relationship_label}' edges!")
 
 
+async def test_schema_only_migration():
+    # 1. Setup test DB and extract schema
+    migration_engine = await setup_test_db()
+    schema = await migration_engine.extract_schema()
+
+    # 2. Setup graph engine
+    graph_engine = await get_graph_engine()
+
+    # 4. Migrate schema only
+    await migrate_relational_database(graph_engine, schema=schema, schema_only=True)
+
+    # 5. Verify number of tables through search
+    search_results = await cognee.search(
+        query_text="How many tables are there in this database",
+        query_type=cognee.SearchType.GRAPH_COMPLETION,
+        top_k=30,
+    )
+    assert any("11" in r for r in search_results), (
+        "Number of tables in the database reported in search_results is either None or not equal to 11"
+    )
+
+    graph_db_provider = os.getenv("GRAPH_DATABASE_PROVIDER", "networkx").lower()
+
+    edge_counts = {
+        "is_part_of": 0,
+        "has_relationship": 0,
+        "foreign_key": 0,
+    }
+
+    if graph_db_provider == "neo4j":
+        for rel_type in edge_counts.keys():
+            query_str = f"""
+            MATCH ()-[r:{rel_type}]->()
+            RETURN count(r) as c
+            """
+            rows = await graph_engine.query(query_str)
+            edge_counts[rel_type] = rows[0]["c"]
+
+    elif graph_db_provider == "kuzu":
+        for rel_type in edge_counts.keys():
+            query_str = f"""
+            MATCH ()-[r:EDGE]->()
+            WHERE r.relationship_name = '{rel_type}'
+            RETURN count(r) as c
+            """
+            rows = await graph_engine.query(query_str)
+            edge_counts[rel_type] = rows[0][0]
+
+    elif graph_db_provider == "networkx":
+        nodes, edges = await graph_engine.get_graph_data()
+        for _, _, key, _ in edges:
+            if key in edge_counts:
+                edge_counts[key] += 1
+
+    else:
+        raise ValueError(f"Unsupported graph database provider: {graph_db_provider}")
+
+    # 7. Assert counts match expected values
+    expected_counts = {
+        "is_part_of": 11,
+        "has_relationship": 22,
+        "foreign_key": 11,
+    }
+
+    for rel_type, expected in expected_counts.items():
+        actual = edge_counts[rel_type]
+        assert actual == expected, (
+            f"Expected {expected} edges for relationship '{rel_type}', but found {actual}"
+        )
+
+    print("Schema-only migration edge counts validated successfully!")
+    print(f"Edge counts: {edge_counts}")
+
+
 async def test_migration_sqlite():
     database_to_migrate_path = os.path.join(pathlib.Path(__file__).parent, "test_data/")
 
@@ -209,6 +283,7 @@ async def test_migration_sqlite():
     )
 
     await relational_db_migration()
+    await test_schema_only_migration()
 
 
 async def test_migration_postgres():
@@ -224,6 +299,7 @@ async def test_migration_postgres():
         }
     )
     await relational_db_migration()
+    await test_schema_only_migration()
 
 
 async def main():

cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py

@@ -1,5 +1,5 @@
 import pytest
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, patch, MagicMock
 from cognee.tasks.storage.index_graph_edges import index_graph_edges
 
 
@@ -16,6 +16,7 @@ async def test_index_graph_edges_success():
         ],
     )
     mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine.get_batch_size = MagicMock(return_value=100)
 
     # Patch the globals of the function so that when it does:
     #   vector_engine = get_vector_engine()