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

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

@@ -1,12 +1,14 @@
 import pytest
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, patch, MagicMock
 
 from cognee.modules.retrieval.utils.brute_force_triplet_search import (
     brute_force_triplet_search,
     get_memory_fragment,
+    format_triplets,
 )
 from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
 from cognee.modules.graph.exceptions.exceptions import EntityNotFoundError
+from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError
 
 
 class MockScoredResult:
@@ -354,20 +356,30 @@ async def test_brute_force_triplet_search_passes_top_k_to_importance_calculation
 
 @pytest.mark.asyncio
 async def test_get_memory_fragment_returns_empty_graph_on_entity_not_found():
-    """Test that get_memory_fragment returns empty graph when entity not found."""
+    """Test that get_memory_fragment returns empty graph when entity not found (line 85)."""
     mock_graph_engine = AsyncMock()
-    mock_graph_engine.project_graph_from_db = AsyncMock(
+
+    # Create a mock fragment that will raise EntityNotFoundError when project_graph_from_db is called
+    mock_fragment = MagicMock(spec=CogneeGraph)
+    mock_fragment.project_graph_from_db = AsyncMock(
         side_effect=EntityNotFoundError("Entity not found")
     )
 
-    with patch(
-        "cognee.modules.retrieval.utils.brute_force_triplet_search.get_graph_engine",
-        return_value=mock_graph_engine,
+    with (
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_graph_engine",
+            return_value=mock_graph_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.CogneeGraph",
+            return_value=mock_fragment,
+        ),
     ):
-        fragment = await get_memory_fragment()
+        result = await get_memory_fragment()
 
-        assert isinstance(fragment, CogneeGraph)
-        assert len(fragment.nodes) == 0
+        # Fragment should be returned even though EntityNotFoundError was raised (pass statement on line 85)
+        assert result == mock_fragment
+        mock_fragment.project_graph_from_db.assert_awaited_once()
 
 
 @pytest.mark.asyncio
@@ -606,3 +618,200 @@ async def test_brute_force_triplet_search_mixed_empty_collections():
 
         call_kwargs = mock_get_fragment_fn.call_args[1]
         assert set(call_kwargs["relevant_ids_to_filter"]) == {"node1", "node2"}
+
+
+def test_format_triplets():
+    """Test format_triplets function."""
+    mock_edge = MagicMock()
+    mock_node1 = MagicMock()
+    mock_node2 = MagicMock()
+
+    mock_node1.attributes = {"name": "Node1", "type": "Entity", "id": "n1"}
+    mock_node2.attributes = {"name": "Node2", "type": "Entity", "id": "n2"}
+    mock_edge.attributes = {"relationship_name": "relates_to", "edge_text": "connects"}
+
+    mock_edge.node1 = mock_node1
+    mock_edge.node2 = mock_node2
+
+    result = format_triplets([mock_edge])
+
+    assert isinstance(result, str)
+    assert "Node1" in result
+    assert "Node2" in result
+    assert "relates_to" in result
+    assert "connects" in result
+
+
+def test_format_triplets_with_none_values():
+    """Test format_triplets filters out None values."""
+    mock_edge = MagicMock()
+    mock_node1 = MagicMock()
+    mock_node2 = MagicMock()
+
+    mock_node1.attributes = {"name": "Node1", "type": None, "id": "n1"}
+    mock_node2.attributes = {"name": "Node2", "type": "Entity", "id": None}
+    mock_edge.attributes = {"relationship_name": "relates_to", "edge_text": None}
+
+    mock_edge.node1 = mock_node1
+    mock_edge.node2 = mock_node2
+
+    result = format_triplets([mock_edge])
+
+    assert "Node1" in result
+    assert "Node2" in result
+    assert "relates_to" in result
+    assert "None" not in result or result.count("None") == 0
+
+
+def test_format_triplets_with_nested_dict():
+    """Test format_triplets handles nested dict attributes (lines 23-35)."""
+    mock_edge = MagicMock()
+    mock_node1 = MagicMock()
+    mock_node2 = MagicMock()
+
+    mock_node1.attributes = {"name": "Node1", "metadata": {"type": "Entity", "id": "n1"}}
+    mock_node2.attributes = {"name": "Node2", "metadata": {"type": "Entity", "id": "n2"}}
+    mock_edge.attributes = {"relationship_name": "relates_to"}
+
+    mock_edge.node1 = mock_node1
+    mock_edge.node2 = mock_node2
+
+    result = format_triplets([mock_edge])
+
+    assert isinstance(result, str)
+    assert "Node1" in result
+    assert "Node2" in result
+    assert "relates_to" in result
+
+
+@pytest.mark.asyncio
+async def test_brute_force_triplet_search_vector_engine_init_error():
+    """Test brute_force_triplet_search handles vector engine initialization error (lines 145-147)."""
+    with (
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_vector_engine"
+        ) as mock_get_vector_engine,
+    ):
+        mock_get_vector_engine.side_effect = Exception("Initialization error")
+
+        with pytest.raises(RuntimeError, match="Initialization error"):
+            await brute_force_triplet_search(query="test query")
+
+
+@pytest.mark.asyncio
+async def test_brute_force_triplet_search_collection_not_found_error():
+    """Test brute_force_triplet_search handles CollectionNotFoundError in search (lines 156-157)."""
+    mock_vector_engine = AsyncMock()
+    mock_embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = mock_embedding_engine
+    mock_embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    mock_vector_engine.search = AsyncMock(
+        side_effect=[
+            CollectionNotFoundError("Collection not found"),
+            [],
+            [],
+        ]
+    )
+
+    with (
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_memory_fragment",
+            return_value=CogneeGraph(),
+        ),
+    ):
+        result = await brute_force_triplet_search(
+            query="test query", collections=["missing_collection", "existing_collection"]
+        )
+
+    assert result == []
+
+
+@pytest.mark.asyncio
+async def test_brute_force_triplet_search_generic_exception():
+    """Test brute_force_triplet_search handles generic exceptions (lines 209-217)."""
+    mock_vector_engine = AsyncMock()
+    mock_embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = mock_embedding_engine
+    mock_embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    mock_vector_engine.search = AsyncMock(side_effect=Exception("Generic error"))
+
+    with (
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+    ):
+        with pytest.raises(Exception, match="Generic error"):
+            await brute_force_triplet_search(query="test query")
+
+
+@pytest.mark.asyncio
+async def test_brute_force_triplet_search_with_node_name_sets_relevant_ids_to_none():
+    """Test brute_force_triplet_search sets relevant_ids_to_filter to None when node_name is provided (line 191)."""
+    mock_vector_engine = AsyncMock()
+    mock_embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = mock_embedding_engine
+    mock_embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    mock_result = MockScoredResult(id="node1", score=0.8, payload={"id": "node1"})
+    mock_vector_engine.search = AsyncMock(return_value=[mock_result])
+
+    mock_fragment = AsyncMock()
+    mock_fragment.map_vector_distances_to_graph_nodes = AsyncMock()
+    mock_fragment.map_vector_distances_to_graph_edges = AsyncMock()
+    mock_fragment.calculate_top_triplet_importances = AsyncMock(return_value=[])
+
+    with (
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_memory_fragment",
+            return_value=mock_fragment,
+        ) as mock_get_fragment,
+    ):
+        await brute_force_triplet_search(query="test query", node_name=["Node1"])
+
+        assert mock_get_fragment.called
+        call_kwargs = mock_get_fragment.call_args.kwargs if mock_get_fragment.call_args else {}
+        assert call_kwargs.get("relevant_ids_to_filter") is None
+
+
+@pytest.mark.asyncio
+async def test_brute_force_triplet_search_collection_not_found_at_top_level():
+    """Test brute_force_triplet_search handles CollectionNotFoundError at top level (line 210)."""
+    mock_vector_engine = AsyncMock()
+    mock_embedding_engine = AsyncMock()
+    mock_vector_engine.embedding_engine = mock_embedding_engine
+    mock_embedding_engine.embed_text = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    mock_result = MockScoredResult(id="node1", score=0.8, payload={"id": "node1"})
+    mock_vector_engine.search = AsyncMock(return_value=[mock_result])
+
+    mock_fragment = AsyncMock()
+    mock_fragment.map_vector_distances_to_graph_nodes = AsyncMock()
+    mock_fragment.map_vector_distances_to_graph_edges = AsyncMock()
+    mock_fragment.calculate_top_triplet_importances = AsyncMock(
+        side_effect=CollectionNotFoundError("Collection not found")
+    )
+
+    with (
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.utils.brute_force_triplet_search.get_memory_fragment",
+            return_value=mock_fragment,
+        ),
+    ):
+        result = await brute_force_triplet_search(query="test query")
+
+    assert result == []

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

@@ -0,0 +1,343 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+from typing import Type
+
+
+class TestGenerateCompletion:
+    @pytest.mark.asyncio
+    async def test_generate_completion_with_system_prompt(self):
+        """Test generate_completion with provided system_prompt."""
+        mock_llm_response = "Generated answer"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.render_prompt",
+                return_value="User prompt text",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import generate_completion
+
+            result = await generate_completion(
+                query="What is AI?",
+                context="AI is artificial intelligence",
+                user_prompt_path="user_prompt.txt",
+                system_prompt_path="system_prompt.txt",
+                system_prompt="Custom system prompt",
+            )
+
+            assert result == mock_llm_response
+            mock_llm.assert_awaited_once_with(
+                text_input="User prompt text",
+                system_prompt="Custom system prompt",
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_generate_completion_without_system_prompt(self):
+        """Test generate_completion reads system_prompt from file when not provided."""
+        mock_llm_response = "Generated answer"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.render_prompt",
+                return_value="User prompt text",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="System prompt from file",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import generate_completion
+
+            result = await generate_completion(
+                query="What is AI?",
+                context="AI is artificial intelligence",
+                user_prompt_path="user_prompt.txt",
+                system_prompt_path="system_prompt.txt",
+            )
+
+            assert result == mock_llm_response
+            mock_llm.assert_awaited_once_with(
+                text_input="User prompt text",
+                system_prompt="System prompt from file",
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_generate_completion_with_conversation_history(self):
+        """Test generate_completion includes conversation_history in system_prompt."""
+        mock_llm_response = "Generated answer"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.render_prompt",
+                return_value="User prompt text",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="System prompt from file",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import generate_completion
+
+            result = await generate_completion(
+                query="What is AI?",
+                context="AI is artificial intelligence",
+                user_prompt_path="user_prompt.txt",
+                system_prompt_path="system_prompt.txt",
+                conversation_history="Previous conversation:\nQ: What is ML?\nA: ML is machine learning",
+            )
+
+            assert result == mock_llm_response
+            expected_system_prompt = (
+                "Previous conversation:\nQ: What is ML?\nA: ML is machine learning"
+                + "\nTASK:"
+                + "System prompt from file"
+            )
+            mock_llm.assert_awaited_once_with(
+                text_input="User prompt text",
+                system_prompt=expected_system_prompt,
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_generate_completion_with_conversation_history_and_custom_system_prompt(self):
+        """Test generate_completion includes conversation_history with custom system_prompt."""
+        mock_llm_response = "Generated answer"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.render_prompt",
+                return_value="User prompt text",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import generate_completion
+
+            result = await generate_completion(
+                query="What is AI?",
+                context="AI is artificial intelligence",
+                user_prompt_path="user_prompt.txt",
+                system_prompt_path="system_prompt.txt",
+                system_prompt="Custom system prompt",
+                conversation_history="Previous conversation:\nQ: What is ML?\nA: ML is machine learning",
+            )
+
+            assert result == mock_llm_response
+            expected_system_prompt = (
+                "Previous conversation:\nQ: What is ML?\nA: ML is machine learning"
+                + "\nTASK:"
+                + "Custom system prompt"
+            )
+            mock_llm.assert_awaited_once_with(
+                text_input="User prompt text",
+                system_prompt=expected_system_prompt,
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_generate_completion_with_response_model(self):
+        """Test generate_completion with custom response_model."""
+        mock_response_model = MagicMock()
+        mock_llm_response = {"answer": "Generated answer"}
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.render_prompt",
+                return_value="User prompt text",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="System prompt from file",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import generate_completion
+
+            result = await generate_completion(
+                query="What is AI?",
+                context="AI is artificial intelligence",
+                user_prompt_path="user_prompt.txt",
+                system_prompt_path="system_prompt.txt",
+                response_model=mock_response_model,
+            )
+
+            assert result == mock_llm_response
+            mock_llm.assert_awaited_once_with(
+                text_input="User prompt text",
+                system_prompt="System prompt from file",
+                response_model=mock_response_model,
+            )
+
+    @pytest.mark.asyncio
+    async def test_generate_completion_render_prompt_args(self):
+        """Test generate_completion passes correct args to render_prompt."""
+        mock_llm_response = "Generated answer"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.render_prompt",
+                return_value="User prompt text",
+            ) as mock_render,
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="System prompt from file",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ),
+        ):
+            from cognee.modules.retrieval.utils.completion import generate_completion
+
+            await generate_completion(
+                query="What is AI?",
+                context="AI is artificial intelligence",
+                user_prompt_path="user_prompt.txt",
+                system_prompt_path="system_prompt.txt",
+            )
+
+            mock_render.assert_called_once_with(
+                "user_prompt.txt",
+                {"question": "What is AI?", "context": "AI is artificial intelligence"},
+            )
+
+
+class TestSummarizeText:
+    @pytest.mark.asyncio
+    async def test_summarize_text_with_system_prompt(self):
+        """Test summarize_text with provided system_prompt."""
+        mock_llm_response = "Summary text"
+
+        with patch(
+            "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+            new_callable=AsyncMock,
+            return_value=mock_llm_response,
+        ) as mock_llm:
+            from cognee.modules.retrieval.utils.completion import summarize_text
+
+            result = await summarize_text(
+                text="Long text to summarize",
+                system_prompt_path="summarize_search_results.txt",
+                system_prompt="Custom summary prompt",
+            )
+
+            assert result == mock_llm_response
+            mock_llm.assert_awaited_once_with(
+                text_input="Long text to summarize",
+                system_prompt="Custom summary prompt",
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_summarize_text_without_system_prompt(self):
+        """Test summarize_text reads system_prompt from file when not provided."""
+        mock_llm_response = "Summary text"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="System prompt from file",
+            ),
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import summarize_text
+
+            result = await summarize_text(
+                text="Long text to summarize",
+                system_prompt_path="summarize_search_results.txt",
+            )
+
+            assert result == mock_llm_response
+            mock_llm.assert_awaited_once_with(
+                text_input="Long text to summarize",
+                system_prompt="System prompt from file",
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_summarize_text_default_prompt_path(self):
+        """Test summarize_text uses default prompt path when not provided."""
+        mock_llm_response = "Summary text"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="Default system prompt",
+            ) as mock_read,
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import summarize_text
+
+            result = await summarize_text(text="Long text to summarize")
+
+            assert result == mock_llm_response
+            mock_read.assert_called_once_with("summarize_search_results.txt")
+            mock_llm.assert_awaited_once_with(
+                text_input="Long text to summarize",
+                system_prompt="Default system prompt",
+                response_model=str,
+            )
+
+    @pytest.mark.asyncio
+    async def test_summarize_text_custom_prompt_path(self):
+        """Test summarize_text uses custom prompt path when provided."""
+        mock_llm_response = "Summary text"
+
+        with (
+            patch(
+                "cognee.modules.retrieval.utils.completion.read_query_prompt",
+                return_value="Custom system prompt",
+            ) as mock_read,
+            patch(
+                "cognee.modules.retrieval.utils.completion.LLMGateway.acreate_structured_output",
+                new_callable=AsyncMock,
+                return_value=mock_llm_response,
+            ) as mock_llm,
+        ):
+            from cognee.modules.retrieval.utils.completion import summarize_text
+
+            result = await summarize_text(
+                text="Long text to summarize",
+                system_prompt_path="custom_prompt.txt",
+            )
+
+            assert result == mock_llm_response
+            mock_read.assert_called_once_with("custom_prompt.txt")
+            mock_llm.assert_awaited_once_with(
+                text_input="Long text to summarize",
+                system_prompt="Custom system prompt",
+                response_model=str,
+            )