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

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

@@ -1,177 +1,469 @@
-import os
 import pytest
-import pathlib
-from typing import Optional, Union
+from unittest.mock import AsyncMock, patch, MagicMock
+from uuid import UUID
 
-import cognee
-from cognee.low_level import setup, DataPoint
-from cognee.tasks.storage import add_data_points
-from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
 )
-
-
-class TestGraphCompletionWithContextExtensionRetriever:
-    @pytest.mark.asyncio
-    async def test_graph_completion_extension_context_simple(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".cognee_system/test_graph_completion_extension_context_simple",
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".data_storage/test_graph_completion_extension_context_simple",
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+@pytest.fixture
+def mock_edge():
+    """Create a mock edge."""
+    edge = MagicMock(spec=Edge)
+    return edge
+
+
+@pytest.mark.asyncio
+async def test_get_triplets_inherited(mock_edge):
+    """Test that get_triplets is inherited from parent class."""
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
+        return_value=[mock_edge],
+    ):
+        triplets = await retriever.get_triplets("test query")
+
+    assert len(triplets) == 1
+    assert triplets[0] == mock_edge
+
+
+@pytest.mark.asyncio
+async def test_init_defaults():
+    """Test GraphCompletionContextExtensionRetriever initialization with defaults."""
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    assert retriever.top_k == 5
+    assert retriever.user_prompt_path == "graph_context_for_question.txt"
+    assert retriever.system_prompt_path == "answer_simple_question.txt"
+
+
+@pytest.mark.asyncio
+async def test_init_custom_params():
+    """Test GraphCompletionContextExtensionRetriever initialization with custom parameters."""
+    retriever = GraphCompletionContextExtensionRetriever(
+        top_k=10,
+        user_prompt_path="custom_user.txt",
+        system_prompt_path="custom_system.txt",
+        system_prompt="Custom prompt",
+        node_type=str,
+        node_name=["node1"],
+        save_interaction=True,
+        wide_search_top_k=200,
+        triplet_distance_penalty=5.0,
+    )
+
+    assert retriever.top_k == 10
+    assert retriever.user_prompt_path == "custom_user.txt"
+    assert retriever.system_prompt_path == "custom_system.txt"
+    assert retriever.system_prompt == "Custom prompt"
+    assert retriever.node_type is str
+    assert retriever.node_name == ["node1"]
+    assert retriever.save_interaction is True
+    assert retriever.wide_search_top_k == 200
+    assert retriever.triplet_distance_penalty == 5.0
+
+
+@pytest.mark.asyncio
+async def test_get_completion_without_context(mock_edge):
+    """Test get_completion retrieves context when not provided."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
+            return_value=[mock_edge],
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion("test query", context_extension_rounds=1)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_provided_context(mock_edge):
+    """Test get_completion uses provided context."""
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion(
+            "test query", context=[mock_edge], context_extension_rounds=1
         )
-        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")
-        company2 = Company(name="Canva")
-        person1 = Person(name="Steve Rodger", works_for=company1)
-        person2 = Person(name="Ike Loma", works_for=company1)
-        person3 = Person(name="Jason Statham", works_for=company1)
-        person4 = Person(name="Mike Broski", works_for=company2)
-        person5 = Person(name="Christina Mayer", works_for=company2)
-
-        entities = [company1, company2, person1, person2, person3, person4, person5]
-
-        await add_data_points(entities)
-
-        retriever = GraphCompletionContextExtensionRetriever()
-
-        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
-        assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
-        assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
-
-        answer = await retriever.get_completion("Who works at Canva?")
-
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_context_extension_rounds(mock_edge):
+    """Test get_completion with multiple context extension rounds."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    # Create a second edge for extension rounds
+    mock_edge2 = MagicMock(spec=Edge)
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch.object(
+            retriever,
+            "get_context",
+            new_callable=AsyncMock,
+            side_effect=[[mock_edge], [mock_edge2]],
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            side_effect=["Resolved context", "Extended context"],  # Different contexts
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            side_effect=[
+                "Extension query",
+                "Generated answer",
+            ],  # Query for extension, then final answer
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion("test query", context_extension_rounds=1)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_context_extension_stops_early(mock_edge):
+    """Test get_completion stops early when no new triplets found."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with (
+        patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[mock_edge]),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            side_effect=[
+                "Extension query",
+                "Generated answer",
+            ],
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        # When get_context returns same triplets, the loop should stop early
+        completion = await retriever.get_completion(
+            "test query", context=[mock_edge], context_extension_rounds=4
         )
 
-    @pytest.mark.asyncio
-    async def test_graph_completion_extension_context_complex(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".cognee_system/test_graph_completion_extension_context_complex",
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".data_storage/test_graph_completion_extension_context_complex",
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_session(mock_edge):
+    """Test get_completion with session caching enabled."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    mock_user = MagicMock()
+    mock_user.id = "test-user-id"
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[mock_edge]),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.get_conversation_history",
+            return_value="Previous conversation",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.summarize_text",
+            return_value="Context summary",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            side_effect=[
+                "Extension query",
+                "Generated answer",
+            ],  # Extension query, then final answer
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.save_conversation_history",
+        ) as mock_save,
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.session_user"
+        ) as mock_session_user,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = True
+        mock_cache_config.return_value = mock_config
+        mock_session_user.get.return_value = mock_user
+
+        completion = await retriever.get_completion(
+            "test query", session_id="test_session", context_extension_rounds=1
         )
-        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
-            metadata: dict = {"index_fields": ["name"]}
-
-        class Car(DataPoint):
-            brand: str
-            model: str
-            year: int
-
-        class Location(DataPoint):
-            country: str
-            city: str
-
-        class Home(DataPoint):
-            location: Location
-            rooms: int
-            sqm: int
-
-        class Person(DataPoint):
-            name: str
-            works_for: Company
-            owns: Optional[list[Union[Car, Home]]] = None
-
-        company1 = Company(name="Figma")
-        company2 = Company(name="Canva")
 
-        person1 = Person(name="Mike Rodger", works_for=company1)
-        person1.owns = [Car(brand="Toyota", model="Camry", year=2020)]
-
-        person2 = Person(name="Ike Loma", works_for=company1)
-        person2.owns = [
-            Car(brand="Tesla", model="Model S", year=2021),
-            Home(location=Location(country="USA", city="New York"), sqm=80, rooms=4),
-        ]
-
-        person3 = Person(name="Jason Statham", works_for=company1)
-
-        person4 = Person(name="Mike Broski", works_for=company2)
-        person4.owns = [Car(brand="Ford", model="Mustang", year=1978)]
-
-        person5 = Person(name="Christina Mayer", works_for=company2)
-        person5.owns = [Car(brand="Honda", model="Civic", year=2023)]
-
-        entities = [company1, company2, person1, person2, person3, person4, person5]
-
-        await add_data_points(entities)
-
-        retriever = GraphCompletionContextExtensionRetriever(top_k=20)
-
-        context = await resolve_edges_to_text(
-            await retriever.get_context("Who works at Figma and drives Tesla?")
-        )
-
-        print(context)
-
-        assert "Mike Rodger --[works_for]--> Figma" in context, "Failed to get Mike Rodger"
-        assert "Ike Loma --[works_for]--> Figma" in context, "Failed to get Ike Loma"
-        assert "Jason Statham --[works_for]--> Figma" in context, "Failed to get Jason Statham"
-
-        answer = await retriever.get_completion("Who works at Figma?")
-
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+    mock_save.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_save_interaction(mock_edge):
+    """Test get_completion with save_interaction enabled."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+    mock_graph_engine.add_edges = AsyncMock()
+
+    retriever = GraphCompletionContextExtensionRetriever(save_interaction=True)
+
+    mock_node1 = MagicMock()
+    mock_node2 = MagicMock()
+    mock_edge.node1 = mock_node1
+    mock_edge.node2 = mock_node2
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[mock_edge]),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            side_effect=[
+                "Extension query",
+                "Generated answer",
+            ],  # Extension query, then final answer
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.extract_uuid_from_node",
+            side_effect=[
+                UUID("550e8400-e29b-41d4-a716-446655440000"),
+                UUID("550e8400-e29b-41d4-a716-446655440001"),
+            ],
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.add_data_points",
+        ) as mock_add_data,
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion(
+            "test query", context=[mock_edge], context_extension_rounds=1
         )
 
-    @pytest.mark.asyncio
-    async def test_get_graph_completion_extension_context_on_empty_graph(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".cognee_system/test_get_graph_completion_extension_context_on_empty_graph",
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    mock_add_data.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_response_model(mock_edge):
+    """Test get_completion with custom response model."""
+    from pydantic import BaseModel
+
+    class TestModel(BaseModel):
+        answer: str
+
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[mock_edge]),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            side_effect=[
+                "Extension query",
+                TestModel(answer="Test answer"),
+            ],  # Extension query, then final answer
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion(
+            "test query", response_model=TestModel, context_extension_rounds=1
         )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".data_storage/test_get_graph_completion_extension_context_on_empty_graph",
-        )
-        cognee.config.data_root_directory(data_directory_path)
-
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
 
-        retriever = GraphCompletionContextExtensionRetriever()
-
-        await setup()
-
-        context = await retriever.get_context("Who works at Figma?")
-        assert context == [], "Context should be empty on an empty graph"
-
-        answer = await retriever.get_completion("Who works at Figma?")
-
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert isinstance(completion[0], TestModel)
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_session_no_user_id(mock_edge):
+    """Test get_completion with session config but no user ID."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[mock_edge]),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            side_effect=[
+                "Extension query",
+                "Generated answer",
+            ],  # Extension query, then final answer
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.session_user"
+        ) as mock_session_user,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = True
+        mock_cache_config.return_value = mock_config
+        mock_session_user.get.return_value = None  # No user
+
+        completion = await retriever.get_completion("test query", context_extension_rounds=1)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+
+
+@pytest.mark.asyncio
+async def test_get_completion_zero_extension_rounds(mock_edge):
+    """Test get_completion with zero context extension rounds."""
+    mock_graph_engine = AsyncMock()
+    mock_graph_engine.is_empty = AsyncMock(return_value=False)
+
+    retriever = GraphCompletionContextExtensionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch.object(retriever, "get_context", new_callable=AsyncMock, return_value=[mock_edge]),
+        patch(
+            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
+            return_value="Resolved context",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch(
+            "cognee.modules.retrieval.graph_completion_context_extension_retriever.CacheConfig"
+        ) as mock_cache_config,
+    ):
+        mock_config = MagicMock()
+        mock_config.caching = False
+        mock_cache_config.return_value = mock_config
+
+        completion = await retriever.get_completion("test query", context_extension_rounds=0)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1