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

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

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

cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py

@@ -2,7 +2,6 @@ import os
 import pytest
 import pathlib
 from typing import Optional, Union
-from pydantic import BaseModel
 
 import cognee
 from cognee.low_level import setup, DataPoint
@@ -11,11 +10,6 @@ from cognee.tasks.storage import add_data_points
 from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
 
 
-class TestAnswer(BaseModel):
-    answer: str
-    explanation: str
-
-
 class TestGraphCompletionCoTRetriever:
     @pytest.mark.asyncio
     async def test_graph_completion_cot_context_simple(self):
@@ -174,48 +168,3 @@ class TestGraphCompletionCoTRetriever:
         assert all(isinstance(item, str) and item.strip() for item in answer), (
             "Answer must contain only non-empty strings"
         )
-
-    @pytest.mark.asyncio
-    async def test_get_structured_completion(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_get_structured_completion"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_get_structured_completion"
-        )
-        cognee.config.data_root_directory(data_directory_path)
-
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-        await setup()
-
-        class Company(DataPoint):
-            name: str
-
-        class Person(DataPoint):
-            name: str
-            works_for: Company
-
-        company1 = Company(name="Figma")
-        person1 = Person(name="Steve Rodger", works_for=company1)
-
-        entities = [company1, person1]
-        await add_data_points(entities)
-
-        retriever = GraphCompletionCotRetriever()
-
-        # Test with string response model (default)
-        string_answer = await retriever.get_structured_completion("Who works at Figma?")
-        assert isinstance(string_answer, str), f"Expected str, got {type(string_answer).__name__}"
-        assert string_answer.strip(), "Answer should not be empty"
-
-        # Test with structured response model
-        structured_answer = await retriever.get_structured_completion(
-            "Who works at Figma?", response_model=TestAnswer
-        )
-        assert isinstance(structured_answer, TestAnswer), (
-            f"Expected TestAnswer, got {type(structured_answer).__name__}"
-        )
-        assert structured_answer.answer.strip(), "Answer field should not be empty"
-        assert structured_answer.explanation.strip(), "Explanation field should not be empty"

cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py

@@ -3,6 +3,7 @@ from typing import List
 import pytest
 import pathlib
 import cognee
+
 from cognee.low_level import setup
 from cognee.tasks.storage import add_data_points
 from cognee.infrastructure.databases.vector import get_vector_engine

cognee/tests/unit/modules/retrieval/structured_output_test.py

@@ -0,0 +1,204 @@
+import asyncio
+
+import pytest
+import cognee
+import pathlib
+import os
+
+from pydantic import BaseModel
+from cognee.low_level import setup, DataPoint
+from cognee.tasks.storage import add_data_points
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.modules.data.processing.document_types import TextDocument
+from cognee.modules.engine.models import Entity, EntityType
+from cognee.modules.retrieval.entity_extractors.DummyEntityExtractor import DummyEntityExtractor
+from cognee.modules.retrieval.context_providers.DummyContextProvider import DummyContextProvider
+from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
+from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
+from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
+    GraphCompletionContextExtensionRetriever,
+)
+from cognee.modules.retrieval.EntityCompletionRetriever import EntityCompletionRetriever
+from cognee.modules.retrieval.temporal_retriever import TemporalRetriever
+from cognee.modules.retrieval.completion_retriever import CompletionRetriever
+
+
+class TestAnswer(BaseModel):
+    answer: str
+    explanation: str
+
+
+def _assert_string_answer(answer: list[str]):
+    assert isinstance(answer, list), f"Expected str, got {type(answer).__name__}"
+    assert all(isinstance(item, str) and item.strip() for item in answer), "Items should be strings"
+    assert all(item.strip() for item in answer), "Items should not be empty"
+
+
+def _assert_structured_answer(answer: list[TestAnswer]):
+    assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+    assert all(isinstance(x, TestAnswer) for x in answer), "Items should be TestAnswer"
+    assert all(x.answer.strip() for x in answer), "Answer text should not be empty"
+    assert all(x.explanation.strip() for x in answer), "Explanation should not be empty"
+
+
+async def _test_get_structured_graph_completion_cot():
+    retriever = GraphCompletionCotRetriever()
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion("Who works at Figma?")
+    _assert_string_answer(string_answer)
+
+    # Test with structured response model
+    structured_answer = await retriever.get_completion(
+        "Who works at Figma?", response_model=TestAnswer
+    )
+    _assert_structured_answer(structured_answer)
+
+
+async def _test_get_structured_graph_completion():
+    retriever = GraphCompletionRetriever()
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion("Who works at Figma?")
+    _assert_string_answer(string_answer)
+
+    # Test with structured response model
+    structured_answer = await retriever.get_completion(
+        "Who works at Figma?", response_model=TestAnswer
+    )
+    _assert_structured_answer(structured_answer)
+
+
+async def _test_get_structured_graph_completion_temporal():
+    retriever = TemporalRetriever()
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion("When did Steve start working at Figma?")
+    _assert_string_answer(string_answer)
+
+    # Test with structured response model
+    structured_answer = await retriever.get_completion(
+        "When did Steve start working at Figma??", response_model=TestAnswer
+    )
+    _assert_structured_answer(structured_answer)
+
+
+async def _test_get_structured_graph_completion_rag():
+    retriever = CompletionRetriever()
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion("Where does Steve work?")
+    _assert_string_answer(string_answer)
+
+    # Test with structured response model
+    structured_answer = await retriever.get_completion(
+        "Where does Steve work?", response_model=TestAnswer
+    )
+    _assert_structured_answer(structured_answer)
+
+
+async def _test_get_structured_graph_completion_context_extension():
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion("Who works at Figma?")
+    _assert_string_answer(string_answer)
+
+    # Test with structured response model
+    structured_answer = await retriever.get_completion(
+        "Who works at Figma?", response_model=TestAnswer
+    )
+    _assert_structured_answer(structured_answer)
+
+
+async def _test_get_structured_entity_completion():
+    retriever = EntityCompletionRetriever(DummyEntityExtractor(), DummyContextProvider())
+
+    # Test with string response model (default)
+    string_answer = await retriever.get_completion("Who is Albert Einstein?")
+    _assert_string_answer(string_answer)
+
+    # Test with structured response model
+    structured_answer = await retriever.get_completion(
+        "Who is Albert Einstein?", response_model=TestAnswer
+    )
+    _assert_structured_answer(structured_answer)
+
+
+class TestStructuredOutputCompletion:
+    @pytest.mark.asyncio
+    async def test_get_structured_completion(self):
+        system_directory_path = os.path.join(
+            pathlib.Path(__file__).parent, ".cognee_system/test_get_structured_completion"
+        )
+        cognee.config.system_root_directory(system_directory_path)
+        data_directory_path = os.path.join(
+            pathlib.Path(__file__).parent, ".data_storage/test_get_structured_completion"
+        )
+        cognee.config.data_root_directory(data_directory_path)
+
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+        await setup()
+
+        class Company(DataPoint):
+            name: str
+
+        class Person(DataPoint):
+            name: str
+            works_for: Company
+            works_since: int
+
+        company1 = Company(name="Figma")
+        person1 = Person(name="Steve Rodger", works_for=company1, works_since=2015)
+
+        entities = [company1, person1]
+        await add_data_points(entities)
+
+        document = TextDocument(
+            name="Steve Rodger's career",
+            raw_data_location="somewhere",
+            external_metadata="",
+            mime_type="text/plain",
+        )
+
+        chunk1 = DocumentChunk(
+            text="Steve Rodger",
+            chunk_size=2,
+            chunk_index=0,
+            cut_type="sentence_end",
+            is_part_of=document,
+            contains=[],
+        )
+        chunk2 = DocumentChunk(
+            text="Mike Broski",
+            chunk_size=2,
+            chunk_index=1,
+            cut_type="sentence_end",
+            is_part_of=document,
+            contains=[],
+        )
+        chunk3 = DocumentChunk(
+            text="Christina Mayer",
+            chunk_size=2,
+            chunk_index=2,
+            cut_type="sentence_end",
+            is_part_of=document,
+            contains=[],
+        )
+
+        entities = [chunk1, chunk2, chunk3]
+        await add_data_points(entities)
+
+        entity_type = EntityType(name="Person", description="A human individual")
+        entity = Entity(name="Albert Einstein", is_a=entity_type, description="A famous physicist")
+
+        entities = [entity]
+        await add_data_points(entities)
+
+        await _test_get_structured_graph_completion_cot()
+        await _test_get_structured_graph_completion()
+        await _test_get_structured_graph_completion_temporal()
+        await _test_get_structured_graph_completion_rag()
+        await _test_get_structured_graph_completion_context_extension()
+        await _test_get_structured_entity_completion()

cognee/tests/unit/modules/retrieval/summaries_retriever_test.py

@@ -13,7 +13,7 @@ from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
 
 
-class TextSummariesRetriever:
+class TestSummariesRetriever:
     @pytest.mark.asyncio
     async def test_chunk_context(self):
         system_directory_path = os.path.join(

cognee/tests/unit/modules/retrieval/temporal_retriever_test.py

@@ -1,4 +1,3 @@
-import asyncio
 from types import SimpleNamespace
 import pytest