cognee 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

cognee/tests/integration/web_url_crawler/test_url_adding_e2e.py

@@ -0,0 +1,344 @@
+import os
+import pytest
+import cognee
+from cognee.infrastructure.files.utils.get_data_file_path import get_data_file_path
+from cognee.infrastructure.loaders.LoaderEngine import LoaderEngine
+from cognee.infrastructure.loaders.external.beautiful_soup_loader import BeautifulSoupLoader
+from cognee.tasks.ingestion import save_data_item_to_storage
+from pathlib import Path
+
+
+@pytest.mark.asyncio
+async def test_url_saves_as_html_file():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        assert file_path.endswith(".html")
+        file = Path(file_path)
+        assert file.exists()
+        assert file.stat().st_size > 0
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")
+
+
+skip_for_tavily = pytest.mark.skipif(
+    os.getenv("TAVILY_API_KEY") is not None,
+    reason="Skipping as Tavily already handles parsing and outputs text",
+)
+
+
+@skip_for_tavily
+@pytest.mark.asyncio
+async def test_saved_html_is_valid():
+    try:
+        from bs4 import BeautifulSoup
+    except ImportError:
+        pytest.fail("Test case requires bs4 installed")
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        content = Path(file_path).read_text()
+
+        soup = BeautifulSoup(content, "html.parser")
+        assert soup.find() is not None, "File should contain parseable HTML"
+
+        has_html_elements = any(
+            [
+                soup.find("html"),
+                soup.find("head"),
+                soup.find("body"),
+                soup.find("div"),
+                soup.find("p"),
+            ]
+        )
+        assert has_html_elements, "File should contain common HTML elements"
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")
+
+
+@pytest.mark.asyncio
+async def test_add_url():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add("https://en.wikipedia.org/wiki/Large_language_model")
+
+
+skip_in_ci = pytest.mark.skipif(
+    os.getenv("GITHUB_ACTIONS") == "true",
+    reason="Skipping in Github for now - before we get TAVILY_API_KEY",
+)
+
+
+@skip_in_ci
+@pytest.mark.asyncio
+async def test_add_url_with_tavily():
+    assert os.getenv("TAVILY_API_KEY") is not None
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add("https://en.wikipedia.org/wiki/Large_language_model")
+
+
+@pytest.mark.asyncio
+async def test_add_url_without_incremental_loading():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    try:
+        await cognee.add(
+            "https://en.wikipedia.org/wiki/Large_language_model",
+            incremental_loading=False,
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to add url: {e}")
+
+
+@pytest.mark.asyncio
+async def test_add_url_with_incremental_loading():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    try:
+        await cognee.add(
+            "https://en.wikipedia.org/wiki/Large_language_model",
+            incremental_loading=True,
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to add url: {e}")
+
+
+@pytest.mark.asyncio
+async def test_add_url_can_define_preferred_loader_as_list_of_str():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add(
+        "https://en.wikipedia.org/wiki/Large_language_model",
+        preferred_loaders=["beautiful_soup_loader"],
+    )
+
+
+@pytest.mark.asyncio
+async def test_add_url_with_extraction_rules():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    extraction_rules = {
+        "title": {"selector": "title"},
+        "headings": {"selector": "h1, h2, h3", "all": True},
+        "links": {"selector": "a", "attr": "href", "all": True},
+        "paragraphs": {"selector": "p", "all": True},
+    }
+
+    try:
+        await cognee.add(
+            "https://en.wikipedia.org/wiki/Large_language_model",
+            preferred_loaders={"beautiful_soup_loader": {"extraction_rules": extraction_rules}},
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to add url: {e}")
+
+
+@pytest.mark.asyncio
+async def test_loader_is_none_by_default():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    extraction_rules = {
+        "title": {"selector": "title"},
+        "headings": {"selector": "h1, h2, h3", "all": True},
+        "links": {"selector": "a", "attr": "href", "all": True},
+        "paragraphs": {"selector": "p", "all": True},
+    }
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        assert file_path.endswith(".html")
+        file = Path(file_path)
+        assert file.exists()
+        assert file.stat().st_size > 0
+
+        loader_engine = LoaderEngine()
+        preferred_loaders = {"beautiful_soup_loader": {"extraction_rules": extraction_rules}}
+        loader = loader_engine.get_loader(
+            file_path,
+            preferred_loaders=preferred_loaders,
+        )
+
+        assert loader is None
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")
+
+
+@pytest.mark.asyncio
+async def test_beautiful_soup_loader_is_selected_loader_if_preferred_loader_provided():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    extraction_rules = {
+        "title": {"selector": "title"},
+        "headings": {"selector": "h1, h2, h3", "all": True},
+        "links": {"selector": "a", "attr": "href", "all": True},
+        "paragraphs": {"selector": "p", "all": True},
+    }
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        assert file_path.endswith(".html")
+        file = Path(file_path)
+        assert file.exists()
+        assert file.stat().st_size > 0
+
+        loader_engine = LoaderEngine()
+        bs_loader = BeautifulSoupLoader()
+        loader_engine.register_loader(bs_loader)
+        preferred_loaders = {"beautiful_soup_loader": {"extraction_rules": extraction_rules}}
+        loader = loader_engine.get_loader(
+            file_path,
+            preferred_loaders=preferred_loaders,
+        )
+
+        assert loader == bs_loader
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")
+
+
+@pytest.mark.asyncio
+async def test_beautiful_soup_loader_works_with_and_without_arguments():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        assert file_path.endswith(".html")
+        file = Path(file_path)
+        assert file.exists()
+        assert file.stat().st_size > 0
+
+        loader_engine = LoaderEngine()
+        bs_loader = BeautifulSoupLoader()
+        loader_engine.register_loader(bs_loader)
+        preferred_loaders = {"beautiful_soup_loader": {}}
+        await loader_engine.load_file(
+            file_path,
+            preferred_loaders=preferred_loaders,
+        )
+        extraction_rules = {
+            "title": {"selector": "title"},
+            "headings": {"selector": "h1, h2, h3", "all": True},
+            "links": {"selector": "a", "attr": "href", "all": True},
+            "paragraphs": {"selector": "p", "all": True},
+        }
+        preferred_loaders = {"beautiful_soup_loader": {"extraction_rules": extraction_rules}}
+        await loader_engine.load_file(
+            file_path,
+            preferred_loaders=preferred_loaders,
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")
+
+
+@pytest.mark.asyncio
+async def test_beautiful_soup_loader_successfully_loads_file_if_required_args_present():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        assert file_path.endswith(".html")
+        file = Path(file_path)
+        assert file.exists()
+        assert file.stat().st_size > 0
+
+        loader_engine = LoaderEngine()
+        bs_loader = BeautifulSoupLoader()
+        loader_engine.register_loader(bs_loader)
+        extraction_rules = {
+            "title": {"selector": "title"},
+            "headings": {"selector": "h1, h2, h3", "all": True},
+            "links": {"selector": "a", "attr": "href", "all": True},
+            "paragraphs": {"selector": "p", "all": True},
+        }
+        preferred_loaders = {"beautiful_soup_loader": {"extraction_rules": extraction_rules}}
+        await loader_engine.load_file(
+            file_path,
+            preferred_loaders=preferred_loaders,
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")
+
+
+@pytest.mark.asyncio
+async def test_beautiful_soup_loads_file_successfully():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    extraction_rules = {
+        "title": {"selector": "title"},
+        "headings": {"selector": "h1, h2, h3", "all": True},
+        "links": {"selector": "a", "attr": "href", "all": True},
+        "paragraphs": {"selector": "p", "all": True},
+    }
+
+    try:
+        original_file_path = await save_data_item_to_storage(
+            "https://en.wikipedia.org/wiki/Large_language_model"
+        )
+        file_path = get_data_file_path(original_file_path)
+        assert file_path.endswith(".html")
+        original_file = Path(file_path)
+        assert original_file.exists()
+        assert original_file.stat().st_size > 0
+
+        loader_engine = LoaderEngine()
+        bs_loader = BeautifulSoupLoader()
+        loader_engine.register_loader(bs_loader)
+        preferred_loaders = {"beautiful_soup_loader": {"extraction_rules": extraction_rules}}
+        loader = loader_engine.get_loader(
+            file_path,
+            preferred_loaders=preferred_loaders,
+        )
+
+        assert loader == bs_loader
+
+        cognee_loaded_txt_path = await loader_engine.load_file(
+            file_path=file_path, preferred_loaders=preferred_loaders
+        )
+
+        cognee_loaded_txt_path = get_data_file_path(cognee_loaded_txt_path)
+
+        assert cognee_loaded_txt_path.endswith(".txt")
+
+        extracted_file = Path(cognee_loaded_txt_path)
+
+        assert extracted_file.exists()
+        assert extracted_file.stat().st_size > 0
+
+        original_basename = original_file.stem
+        extracted_basename = extracted_file.stem
+        assert original_basename == extracted_basename, (
+            f"Expected same base name: {original_basename} vs {extracted_basename}"
+        )
+    except Exception as e:
+        pytest.fail(f"Failed to save data item to storage: {e}")

cognee/tests/subprocesses/reader.py

@@ -0,0 +1,25 @@
+import asyncio
+import time
+from cognee.infrastructure.databases.graph.kuzu.adapter import KuzuAdapter
+
+# This will create the test.db if it doesn't exist
+
+
+async def main():
+    adapter = KuzuAdapter("test.db")
+    result = await adapter.query("MATCH (n:Node) RETURN COUNT(n)")
+    print(f"Reader: Found {result[0][0]} nodes")
+    result = await adapter.query("MATCH (n:Node) RETURN COUNT(n)")
+    print(f"Reader: Found {result[0][0]} nodes")
+    result = await adapter.query("MATCH (n:Node) RETURN COUNT(n)")
+    print(f"Reader: Found {result[0][0]} nodes")
+    result = await adapter.query("MATCH (n:Node) RETURN COUNT(n)")
+    print(f"Reader: Found {result[0][0]} nodes")
+    result = await adapter.query("MATCH (n:Node) RETURN COUNT(n)")
+    print(f"Reader: Found {result} nodes")
+    result = await adapter.query("MATCH (n:Node) RETURN COUNT(n)")
+    print(f"Reader: Found {result[0][0]} nodes")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

cognee/tests/subprocesses/simple_cognify_1.py

@@ -0,0 +1,31 @@
+import asyncio
+import cognee
+from cognee.shared.logging_utils import setup_logging, INFO
+from cognee.api.v1.search import SearchType
+
+
+async def main():
+    await cognee.cognify(datasets=["first_cognify_dataset"])
+
+    query_text = (
+        "Tell me what is in the context. Additionally write out 'FIRST_COGNIFY' before your answer"
+    )
+    search_results = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION,
+        query_text=query_text,
+        datasets=["first_cognify_dataset"],
+    )
+
+    print("Search results:")
+    for result_text in search_results:
+        print(result_text)
+
+
+if __name__ == "__main__":
+    logger = setup_logging(log_level=INFO)
+    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/subprocesses/simple_cognify_2.py

@@ -0,0 +1,31 @@
+import asyncio
+import cognee
+from cognee.shared.logging_utils import setup_logging, INFO
+from cognee.api.v1.search import SearchType
+
+
+async def main():
+    await cognee.cognify(datasets=["second_cognify_dataset"])
+
+    query_text = (
+        "Tell me what is in the context. Additionally write out 'SECOND_COGNIFY' before your answer"
+    )
+    search_results = await cognee.search(
+        query_type=SearchType.GRAPH_COMPLETION,
+        query_text=query_text,
+        datasets=["second_cognify_dataset"],
+    )
+
+    print("Search results:")
+    for result_text in search_results:
+        print(result_text)
+
+
+if __name__ == "__main__":
+    logger = setup_logging(log_level=INFO)
+    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/subprocesses/writer.py

@@ -0,0 +1,32 @@
+import asyncio
+import time
+import uuid
+from cognee.modules.data.processing.document_types import PdfDocument
+from cognee.infrastructure.databases.graph.kuzu.adapter import KuzuAdapter
+
+
+def create_node(name):
+    document = PdfDocument(
+        id=uuid.uuid4(),
+        name=name,
+        raw_data_location=name,
+        external_metadata="test_external_metadata",
+        mime_type="test_mime",
+    )
+    return document
+
+
+async def main():
+    adapter = KuzuAdapter("test.db")
+    nodes = [create_node(f"Node{i}") for i in range(5)]
+
+    print("Writer: Starting...")
+    await adapter.add_nodes(nodes)
+
+    print("writer finished...")
+
+    time.sleep(10)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

cognee/tests/tasks/descriptive_metrics/metrics_test_utils.py

@@ -1,7 +1,6 @@
 from typing import List
 from cognee.infrastructure.engine import DataPoint
 from cognee.tasks.storage.add_data_points import add_data_points
-from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine
 import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
 import json
@@ -64,7 +63,6 @@ async def create_connected_test_graph():
 
 
 async def get_metrics(provider: str, include_optional=True):
-    create_graph_engine.cache_clear()
     cognee.config.set_graph_database_provider(provider)
     graph_engine = await get_graph_engine()
     await graph_engine.delete_graph()

cognee/tests/tasks/descriptive_metrics/neo4j_metrics_test.py

@@ -1,7 +1,12 @@
-from cognee.tests.tasks.descriptive_metrics.metrics_test_utils import assert_metrics
 import asyncio
 
 
+async def main():
+    from cognee.tests.tasks.descriptive_metrics.metrics_test_utils import assert_metrics
+
+    await assert_metrics(provider="neo4j", include_optional=False)
+    await assert_metrics(provider="neo4j", include_optional=True)
+
+
 if __name__ == "__main__":
-    asyncio.run(assert_metrics(provider="neo4j", include_optional=False))
-    asyncio.run(assert_metrics(provider="neo4j", include_optional=True))
+    asyncio.run(main())

cognee/tests/tasks/entity_extraction/entity_extraction_test.py

@@ -0,0 +1,89 @@
+import os
+import pathlib
+import asyncio
+
+import cognee
+import cognee.modules.ingestion as ingestion
+from cognee.infrastructure.llm import get_max_chunk_tokens
+from cognee.infrastructure.llm.extraction import extract_content_graph
+from cognee.modules.chunking.TextChunker import TextChunker
+from cognee.modules.data.processing.document_types import TextDocument
+from cognee.modules.users.methods import get_default_user
+from cognee.shared.data_models import KnowledgeGraph
+from cognee.tasks.documents import extract_chunks_from_documents
+from cognee.tasks.ingestion import save_data_item_to_storage
+from cognee.infrastructure.files.utils.open_data_file import open_data_file
+
+
+async def extract_graphs(document_chunks):
+    """
+    Extract graph, and check if entities are present
+    """
+
+    extraction_results = await asyncio.gather(
+        *[extract_content_graph(chunk.text, KnowledgeGraph) for chunk in document_chunks]
+    )
+
+    return all(
+        any(
+            term in node.name.lower()
+            for extraction_result in extraction_results
+            for node in extraction_result.nodes
+        )
+        for term in ("qubit", "algorithm", "superposition")
+    )
+
+
+async def main():
+    """
+    Test how well the entity extraction works. Repeat graph generation a few times.
+    If 80% or more graphs are correctly generated, the test passes.
+    """
+
+    file_path = os.path.join(
+        pathlib.Path(__file__).parent.parent.parent, "test_data/Quantum_computers.txt"
+    )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add("NLP is a subfield of computer science.")
+
+    original_file_path = await save_data_item_to_storage(file_path)
+
+    async with open_data_file(original_file_path) as file:
+        classified_data = ingestion.classify(file)
+
+        # data_id is the hash of original file contents + owner id to avoid duplicate data
+        data_id = ingestion.identify(classified_data, await get_default_user())
+
+    await cognee.add(file_path)
+
+    text_document = TextDocument(
+        id=data_id,
+        type="text",
+        mime_type="text/plain",
+        name="quantum_text",
+        raw_data_location=file_path,
+        external_metadata=None,
+    )
+
+    document_chunks = []
+    async for chunk in extract_chunks_from_documents(
+        [text_document], max_chunk_size=get_max_chunk_tokens(), chunker=TextChunker
+    ):
+        document_chunks.append(chunk)
+
+    number_of_reps = 5
+
+    graph_results = await asyncio.gather(
+        *[extract_graphs(document_chunks) for _ in range(number_of_reps)]
+    )
+
+    correct_graphs = [result for result in graph_results if result]
+
+    assert len(correct_graphs) >= 0.8 * number_of_reps
+
+
+if __name__ == "__main__":
+    asyncio.run(main())