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

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,
+            )

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

@@ -0,0 +1,157 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+
+from cognee.modules.retrieval.graph_summary_completion_retriever import (
+    GraphSummaryCompletionRetriever,
+)
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+@pytest.fixture
+def mock_edge():
+    """Create a mock edge."""
+    edge = MagicMock(spec=Edge)
+    return edge
+
+
+class TestGraphSummaryCompletionRetriever:
+    @pytest.mark.asyncio
+    async def test_init_defaults(self):
+        """Test GraphSummaryCompletionRetriever initialization with defaults."""
+        retriever = GraphSummaryCompletionRetriever()
+
+        assert retriever.summarize_prompt_path == "summarize_search_results.txt"
+        assert retriever.user_prompt_path == "graph_context_for_question.txt"
+        assert retriever.system_prompt_path == "answer_simple_question.txt"
+        assert retriever.top_k == 5
+        assert retriever.save_interaction is False
+
+    @pytest.mark.asyncio
+    async def test_init_custom_params(self):
+        """Test GraphSummaryCompletionRetriever initialization with custom parameters."""
+        retriever = GraphSummaryCompletionRetriever(
+            user_prompt_path="custom_user.txt",
+            system_prompt_path="custom_system.txt",
+            summarize_prompt_path="custom_summarize.txt",
+            system_prompt="Custom system prompt",
+            top_k=10,
+            save_interaction=True,
+            wide_search_top_k=200,
+            triplet_distance_penalty=2.5,
+        )
+
+        assert retriever.summarize_prompt_path == "custom_summarize.txt"
+        assert retriever.user_prompt_path == "custom_user.txt"
+        assert retriever.system_prompt_path == "custom_system.txt"
+        assert retriever.top_k == 10
+        assert retriever.save_interaction is True
+
+    @pytest.mark.asyncio
+    async def test_resolve_edges_to_text_calls_super_and_summarizes(self, mock_edge):
+        """Test resolve_edges_to_text calls super method and then summarizes."""
+        retriever = GraphSummaryCompletionRetriever(
+            summarize_prompt_path="custom_summarize.txt",
+            system_prompt="Custom system prompt",
+        )
+
+        with (
+            patch(
+                "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text",
+                new_callable=AsyncMock,
+                return_value="Resolved edges text",
+            ) as mock_super_resolve,
+            patch(
+                "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text",
+                new_callable=AsyncMock,
+                return_value="Summarized text",
+            ) as mock_summarize,
+        ):
+            result = await retriever.resolve_edges_to_text([mock_edge])
+
+            assert result == "Summarized text"
+            mock_super_resolve.assert_awaited_once_with([mock_edge])
+            mock_summarize.assert_awaited_once_with(
+                "Resolved edges text",
+                "custom_summarize.txt",
+                "Custom system prompt",
+            )
+
+    @pytest.mark.asyncio
+    async def test_resolve_edges_to_text_with_default_system_prompt(self, mock_edge):
+        """Test resolve_edges_to_text uses None for system_prompt when not provided."""
+        retriever = GraphSummaryCompletionRetriever()
+
+        with (
+            patch(
+                "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text",
+                new_callable=AsyncMock,
+                return_value="Resolved edges text",
+            ),
+            patch(
+                "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text",
+                new_callable=AsyncMock,
+                return_value="Summarized text",
+            ) as mock_summarize,
+        ):
+            await retriever.resolve_edges_to_text([mock_edge])
+
+            mock_summarize.assert_awaited_once_with(
+                "Resolved edges text",
+                "summarize_search_results.txt",
+                None,
+            )
+
+    @pytest.mark.asyncio
+    async def test_resolve_edges_to_text_with_empty_edges(self):
+        """Test resolve_edges_to_text handles empty edges list."""
+        retriever = GraphSummaryCompletionRetriever()
+
+        with (
+            patch(
+                "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text",
+                new_callable=AsyncMock,
+                return_value="",
+            ),
+            patch(
+                "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text",
+                new_callable=AsyncMock,
+                return_value="Empty summary",
+            ) as mock_summarize,
+        ):
+            result = await retriever.resolve_edges_to_text([])
+
+            assert result == "Empty summary"
+            mock_summarize.assert_awaited_once_with(
+                "",
+                "summarize_search_results.txt",
+                None,
+            )
+
+    @pytest.mark.asyncio
+    async def test_resolve_edges_to_text_with_multiple_edges(self, mock_edge):
+        """Test resolve_edges_to_text handles multiple edges."""
+        retriever = GraphSummaryCompletionRetriever()
+
+        mock_edge2 = MagicMock(spec=Edge)
+        mock_edge3 = MagicMock(spec=Edge)
+
+        with (
+            patch(
+                "cognee.modules.retrieval.graph_completion_retriever.GraphCompletionRetriever.resolve_edges_to_text",
+                new_callable=AsyncMock,
+                return_value="Multiple edges resolved text",
+            ),
+            patch(
+                "cognee.modules.retrieval.graph_summary_completion_retriever.summarize_text",
+                new_callable=AsyncMock,
+                return_value="Multiple edges summarized",
+            ) as mock_summarize,
+        ):
+            result = await retriever.resolve_edges_to_text([mock_edge, mock_edge2, mock_edge3])
+
+            assert result == "Multiple edges summarized"
+            mock_summarize.assert_awaited_once_with(
+                "Multiple edges resolved text",
+                "summarize_search_results.txt",
+                None,
+            )