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

cognee/tests/test_pipeline_cache.py

@@ -1,164 +0,0 @@
-"""
-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/unit/infrastructure/llm/test_llm_config.py

@@ -1,46 +0,0 @@
-import pytest
-
-from cognee.infrastructure.llm.config import LLMConfig
-
-
-def test_strip_quotes_from_strings():
-    """
-    Test if the LLMConfig.strip_quotes_from_strings model validator behaves as expected.
-    """
-    config = LLMConfig(
-        # Strings with surrounding double quotes ("value" → value)
-        llm_api_key='"double_value"',
-        # Strings with surrounding single quotes ('value' → value)
-        llm_endpoint="'single_value'",
-        # Strings without quotes (value → value)
-        llm_api_version="no_quotes_value",
-        # Empty quoted strings ("" → empty string)
-        fallback_model='""',
-        # None values (should remain None)
-        baml_llm_api_key=None,
-        # Mixed quotes ("value' → unchanged)
-        fallback_endpoint="\"mixed_quote'",
-        # Strings with internal quotes ("internal\"quotes" → internal"quotes")
-        baml_llm_model='"internal"quotes"',
-    )
-
-    # Strings with surrounding double quotes ("value" → value)
-    assert config.llm_api_key == "double_value"
-
-    # Strings with surrounding single quotes ('value' → value)
-    assert config.llm_endpoint == "single_value"
-
-    # Strings without quotes (value → value)
-    assert config.llm_api_version == "no_quotes_value"
-
-    # Empty quoted strings ("" → empty string)
-    assert config.fallback_model == ""
-
-    # None values (should remain None)
-    assert config.baml_llm_api_key is None
-
-    # Mixed quotes ("value' → unchanged)
-    assert config.fallback_endpoint == "\"mixed_quote'"
-
-    # Strings with internal quotes ("internal\"quotes" → internal"quotes")
-    assert config.baml_llm_model == 'internal"quotes'

cognee/tests/unit/modules/memify_tasks/test_get_triplet_datapoints.py

@@ -1,214 +0,0 @@
-import sys
-import pytest
-from unittest.mock import AsyncMock, patch
-
-from cognee.tasks.memify.get_triplet_datapoints import get_triplet_datapoints
-from cognee.modules.engine.models import Triplet
-from cognee.modules.engine.models.Entity import Entity
-from cognee.infrastructure.engine import DataPoint
-from cognee.modules.graph.models.EdgeType import EdgeType
-
-
-get_triplet_datapoints_module = sys.modules["cognee.tasks.memify.get_triplet_datapoints"]
-
-
-@pytest.fixture
-def mock_graph_engine():
-    """Create a mock graph engine with get_triplets_batch method."""
-    engine = AsyncMock()
-    engine.get_triplets_batch = AsyncMock()
-    return engine
-
-
-@pytest.mark.asyncio
-async def test_get_triplet_datapoints_success(mock_graph_engine):
-    """Test successful extraction of triplet datapoints."""
-    mock_triplets_batch = [
-        {
-            "start_node": {
-                "id": "node1",
-                "type": "Entity",
-                "name": "Alice",
-                "description": "A person",
-            },
-            "end_node": {
-                "id": "node2",
-                "type": "Entity",
-                "name": "Bob",
-                "description": "Another person",
-            },
-            "relationship_properties": {
-                "relationship_name": "knows",
-            },
-        }
-    ]
-
-    mock_graph_engine.get_triplets_batch.return_value = mock_triplets_batch
-
-    with (
-        patch.object(
-            get_triplet_datapoints_module, "get_graph_engine", return_value=mock_graph_engine
-        ),
-        patch.object(get_triplet_datapoints_module, "get_all_subclasses") as mock_get_subclasses,
-    ):
-        mock_get_subclasses.return_value = [Triplet, EdgeType, Entity]
-
-        triplets = []
-        async for triplet in get_triplet_datapoints([{}], triplets_batch_size=100):
-            triplets.append(triplet)
-
-        assert len(triplets) == 1
-        assert isinstance(triplets[0], Triplet)
-        assert triplets[0].from_node_id == "node1"
-        assert triplets[0].to_node_id == "node2"
-        assert "Alice" in triplets[0].text
-        assert "knows" in triplets[0].text
-        assert "Bob" in triplets[0].text
-
-
-@pytest.mark.asyncio
-async def test_get_triplet_datapoints_edge_text_priority_and_fallback(mock_graph_engine):
-    """Test that edge_text is prioritized over relationship_name, and fallback works."""
-
-    class MockEntity(DataPoint):
-        name: str
-        metadata: dict = {"index_fields": ["name"]}
-
-    mock_triplets_batch = [
-        {
-            "start_node": {"id": "node1", "type": "Entity", "name": "Alice"},
-            "end_node": {"id": "node2", "type": "Entity", "name": "Bob"},
-            "relationship_properties": {
-                "relationship_name": "knows",
-                "edge_text": "has a close friendship with",
-            },
-        },
-        {
-            "start_node": {"id": "node3", "type": "Entity", "name": "Charlie"},
-            "end_node": {"id": "node4", "type": "Entity", "name": "Diana"},
-            "relationship_properties": {
-                "relationship_name": "works_with",
-            },
-        },
-    ]
-
-    mock_graph_engine.get_triplets_batch.return_value = mock_triplets_batch
-
-    with (
-        patch.object(
-            get_triplet_datapoints_module, "get_graph_engine", return_value=mock_graph_engine
-        ),
-        patch.object(get_triplet_datapoints_module, "get_all_subclasses") as mock_get_subclasses,
-    ):
-        mock_get_subclasses.return_value = [Triplet, EdgeType, MockEntity]
-
-        triplets = []
-        async for triplet in get_triplet_datapoints([{}], triplets_batch_size=100):
-            triplets.append(triplet)
-
-        assert len(triplets) == 2
-        assert "has a close friendship with" in triplets[0].text
-        assert "knows" not in triplets[0].text
-        assert "works_with" in triplets[1].text
-
-
-@pytest.mark.asyncio
-async def test_get_triplet_datapoints_skips_missing_node_ids(mock_graph_engine):
-    """Test that triplets with missing node IDs are skipped."""
-
-    class MockEntity(DataPoint):
-        name: str
-        metadata: dict = {"index_fields": ["name"]}
-
-    mock_triplets_batch = [
-        {
-            "start_node": {"id": "", "type": "Entity", "name": "Alice"},
-            "end_node": {"id": "node2", "type": "Entity", "name": "Bob"},
-            "relationship_properties": {"relationship_name": "knows"},
-        },
-        {
-            "start_node": {"id": "node3", "type": "Entity", "name": "Charlie"},
-            "end_node": {"id": "node4", "type": "Entity", "name": "Diana"},
-            "relationship_properties": {"relationship_name": "works_with"},
-        },
-    ]
-
-    mock_graph_engine.get_triplets_batch.return_value = mock_triplets_batch
-
-    with (
-        patch.object(
-            get_triplet_datapoints_module, "get_graph_engine", return_value=mock_graph_engine
-        ),
-        patch.object(get_triplet_datapoints_module, "get_all_subclasses") as mock_get_subclasses,
-    ):
-        mock_get_subclasses.return_value = [Triplet, EdgeType, MockEntity]
-
-        triplets = []
-        async for triplet in get_triplet_datapoints([{}], triplets_batch_size=100):
-            triplets.append(triplet)
-
-        assert len(triplets) == 1
-        assert triplets[0].from_node_id == "node3"
-
-
-@pytest.mark.asyncio
-async def test_get_triplet_datapoints_error_handling(mock_graph_engine):
-    """Test that errors are handled correctly - invalid data is skipped, query errors propagate."""
-
-    class MockEntity(DataPoint):
-        name: str
-        metadata: dict = {"index_fields": ["name"]}
-
-    mock_triplets_batch = [
-        {
-            "start_node": {"id": "node1", "type": "Entity", "name": "Alice"},
-            "end_node": {"id": "node2", "type": "Entity", "name": "Bob"},
-            "relationship_properties": {"relationship_name": "knows"},
-        },
-        {
-            "start_node": None,
-            "end_node": {"id": "node4", "type": "Entity", "name": "Diana"},
-            "relationship_properties": {"relationship_name": "works_with"},
-        },
-    ]
-
-    mock_graph_engine.get_triplets_batch.return_value = mock_triplets_batch
-
-    with (
-        patch.object(
-            get_triplet_datapoints_module, "get_graph_engine", return_value=mock_graph_engine
-        ),
-        patch.object(get_triplet_datapoints_module, "get_all_subclasses") as mock_get_subclasses,
-    ):
-        mock_get_subclasses.return_value = [Triplet, EdgeType, MockEntity]
-
-        triplets = []
-        async for triplet in get_triplet_datapoints([{}], triplets_batch_size=100):
-            triplets.append(triplet)
-
-        assert len(triplets) == 1
-        assert triplets[0].from_node_id == "node1"
-
-    mock_graph_engine.get_triplets_batch.side_effect = Exception("Database connection error")
-
-    with patch.object(
-        get_triplet_datapoints_module, "get_graph_engine", return_value=mock_graph_engine
-    ):
-        triplets = []
-        with pytest.raises(Exception, match="Database connection error"):
-            async for triplet in get_triplet_datapoints([{}], triplets_batch_size=100):
-                triplets.append(triplet)
-
-
-@pytest.mark.asyncio
-async def test_get_triplet_datapoints_no_get_triplets_batch_method(mock_graph_engine):
-    """Test that NotImplementedError is raised when graph engine lacks get_triplets_batch."""
-    del mock_graph_engine.get_triplets_batch
-
-    with patch.object(
-        get_triplet_datapoints_module, "get_graph_engine", return_value=mock_graph_engine
-    ):
-        triplets = []
-        with pytest.raises(NotImplementedError, match="does not support get_triplets_batch"):
-            async for triplet in get_triplet_datapoints([{}], triplets_batch_size=100):
-                triplets.append(triplet)