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

cognee/tests/test_edge_centered_payload.py

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

cognee/tests/test_pipeline_cache.py

@@ -0,0 +1,164 @@
+"""
+Test suite for the pipeline_cache feature in Cognee pipelines.
+
+This module tests the behavior of the `pipeline_cache` parameter which controls
+whether a pipeline should skip re-execution when it has already been completed
+for the same dataset.
+
+Architecture Overview:
+---------------------
+The pipeline_cache mechanism works at the dataset level:
+1. When a pipeline runs, it logs its status (INITIATED -> STARTED -> COMPLETED)
+2. Before each run, `check_pipeline_run_qualification()` checks the pipeline status
+3. If `use_pipeline_cache=True` and status is COMPLETED/STARTED, the pipeline skips
+4. If `use_pipeline_cache=False`, the pipeline always re-executes regardless of status
+"""
+
+import pytest
+
+import cognee
+from cognee.modules.pipelines.tasks.task import Task
+from cognee.modules.pipelines import run_pipeline
+from cognee.modules.users.methods import get_default_user
+
+from cognee.modules.pipelines.layers.reset_dataset_pipeline_run_status import (
+    reset_dataset_pipeline_run_status,
+)
+from cognee.infrastructure.databases.relational import create_db_and_tables
+
+
+class ExecutionCounter:
+    """Helper class to track task execution counts."""
+
+    def __init__(self):
+        self.count = 0
+
+
+async def create_counting_task(data, counter: ExecutionCounter):
+    """Create a task that increments a counter from the ExecutionCounter instance when executed."""
+    counter.count += 1
+    return counter
+
+
+class TestPipelineCache:
+    """Tests for basic pipeline_cache on/off behavior."""
+
+    @pytest.mark.asyncio
+    async def test_pipeline_cache_off_allows_reexecution(self):
+        """
+        Test that with use_pipeline_cache=False, the pipeline re-executes
+        even when it has already completed for the dataset.
+
+        Expected behavior:
+        - First run: Pipeline executes fully, task runs once
+        - Second run: Pipeline executes again, task runs again (total: 2 times)
+        """
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+        await create_db_and_tables()
+
+        counter = ExecutionCounter()
+        user = await get_default_user()
+
+        tasks = [Task(create_counting_task, counter=counter)]
+
+        # First run
+        pipeline_results_1 = []
+        async for result in run_pipeline(
+            tasks=tasks,
+            datasets="test_dataset_cache_off",
+            data=["sample data"],  # Data is necessary to trigger processing
+            user=user,
+            pipeline_name="test_cache_off_pipeline",
+            use_pipeline_cache=False,
+        ):
+            pipeline_results_1.append(result)
+
+        first_run_count = counter.count
+        assert first_run_count >= 1, "Task should have executed at least once on first run"
+
+        # Second run with pipeline_cache=False
+        pipeline_results_2 = []
+        async for result in run_pipeline(
+            tasks=tasks,
+            datasets="test_dataset_cache_off",
+            data=["sample data"],  # Data is necessary to trigger processing
+            user=user,
+            pipeline_name="test_cache_off_pipeline",
+            use_pipeline_cache=False,
+        ):
+            pipeline_results_2.append(result)
+
+        second_run_count = counter.count
+        assert second_run_count > first_run_count, (
+            f"With pipeline_cache=False, task should re-execute. "
+            f"First run: {first_run_count}, After second run: {second_run_count}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_reset_pipeline_status_allows_reexecution_with_cache(self):
+        """
+        Test that resetting pipeline status allows re-execution even with
+        pipeline_cache=True.
+        """
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+        await create_db_and_tables()
+
+        counter = ExecutionCounter()
+        user = await get_default_user()
+        dataset_name = "reset_status_test"
+        pipeline_name = "test_reset_pipeline"
+
+        tasks = [Task(create_counting_task, counter=counter)]
+
+        # First run
+        pipeline_result = []
+        async for result in run_pipeline(
+            tasks=tasks,
+            datasets=dataset_name,
+            user=user,
+            data=["sample data"],  # Data is necessary to trigger processing
+            pipeline_name=pipeline_name,
+            use_pipeline_cache=True,
+        ):
+            pipeline_result.append(result)
+
+        first_run_count = counter.count
+        assert first_run_count >= 1
+
+        # Second run without reset - should skip
+        async for _ in run_pipeline(
+            tasks=tasks,
+            datasets=dataset_name,
+            user=user,
+            data=["sample data"],  # Data is necessary to trigger processing
+            pipeline_name=pipeline_name,
+            use_pipeline_cache=True,
+        ):
+            pass
+
+        after_second_run = counter.count
+        assert after_second_run == first_run_count, "Should have skipped due to cache"
+
+        # Reset the pipeline status
+        await reset_dataset_pipeline_run_status(
+            pipeline_result[0].dataset_id, user, pipeline_names=[pipeline_name]
+        )
+
+        # Third run after reset - should execute
+        async for _ in run_pipeline(
+            tasks=tasks,
+            datasets=dataset_name,
+            user=user,
+            data=["sample data"],  # Data is necessary to trigger processing
+            pipeline_name=pipeline_name,
+            use_pipeline_cache=True,
+        ):
+            pass
+
+        after_reset_run = counter.count
+        assert after_reset_run > after_second_run, (
+            f"After reset, pipeline should re-execute. "
+            f"Before reset: {after_second_run}, After reset run: {after_reset_run}"
+        )

cognee/tests/test_search_db.py

@@ -2,6 +2,7 @@ import pathlib
 import os
 import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
@@ -12,8 +13,10 @@ from cognee.modules.retrieval.graph_completion_cot_retriever import GraphComplet
 from cognee.modules.retrieval.graph_summary_completion_retriever import (
     GraphSummaryCompletionRetriever,
 )
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.search.types import SearchType
+from cognee.modules.users.methods import get_default_user
 from collections import Counter
 
 logger = get_logger()
@@ -37,6 +40,23 @@ async def main():
 
     await cognee.cognify([dataset_name])
 
+    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, triplets_batch_size=5)
+
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+
+    vector_engine = get_vector_engine()
+    collection = await vector_engine.search(
+        query_text="Test", limit=None, collection_name="Triplet_text"
+    )
+
+    assert len(edges) == len(collection), (
+        f"Expected {len(edges)} edges but got {len(collection)} in Triplet_text collection"
+    )
+
     context_gk = await GraphCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
@@ -49,6 +69,9 @@ async def main():
     context_gk_sum = await GraphSummaryCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
+    context_triplet = await TripletRetriever().get_context(
+        query="Next to which country is Germany located?"
+    )
 
     for name, context in [
         ("GraphCompletionRetriever", context_gk),
@@ -65,6 +88,13 @@ async def main():
             f"{name}: Context did not contain 'germany' or 'netherlands'; got: {context!r}"
         )
 
+    assert isinstance(context_triplet, str), "TripletRetriever: Context should be a string"
+    assert len(context_triplet) > 0, "TripletRetriever: Context should not be empty"
+    lower_triplet = context_triplet.lower()
+    assert "germany" in lower_triplet or "netherlands" in lower_triplet, (
+        f"TripletRetriever: Context did not contain 'germany' or 'netherlands'; got: {context_triplet!r}"
+    )
+
     triplets_gk = await GraphCompletionRetriever().get_triplets(
         query="Next to which country is Germany located?"
     )
@@ -129,6 +159,11 @@ async def main():
         query_text="Next to which country is Germany located?",
         save_interaction=True,
     )
+    completion_triplet = await cognee.search(
+        query_type=SearchType.TRIPLET_COMPLETION,
+        query_text="Next to which country is Germany located?",
+        save_interaction=True,
+    )
 
     await cognee.search(
         query_type=SearchType.FEEDBACK,
@@ -141,6 +176,7 @@ async def main():
         ("GRAPH_COMPLETION_COT", completion_cot),
         ("GRAPH_COMPLETION_CONTEXT_EXTENSION", completion_ext),
         ("GRAPH_SUMMARY_COMPLETION", completion_sum),
+        ("TRIPLET_COMPLETION", completion_triplet),
     ]:
         assert isinstance(search_results, list), f"{name}: should return a list"
         assert len(search_results) == 1, (
@@ -168,7 +204,7 @@ async def main():
 
     # Assert there are exactly 4 CogneeUserInteraction nodes.
     assert type_counts.get("CogneeUserInteraction", 0) == 4, (
-        f"Expected exactly four DCogneeUserInteraction nodes, but found {type_counts.get('CogneeUserInteraction', 0)}"
+        f"Expected exactly four CogneeUserInteraction nodes, but found {type_counts.get('CogneeUserInteraction', 0)}"
     )
 
     # Assert there is exactly two CogneeUserFeedback nodes.