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

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

@@ -1,205 +1,334 @@
-import os
-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
-from cognee.modules.chunking.models import DocumentChunk
-from cognee.modules.data.processing.document_types import TextDocument
-from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from unittest.mock import AsyncMock, patch, MagicMock
+
 from cognee.modules.retrieval.completion_retriever import CompletionRetriever
-from cognee.infrastructure.engine import DataPoint
-from cognee.modules.data.processing.document_types import Document
-from cognee.modules.engine.models import Entity
+from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
 
 
-class DocumentChunkWithEntities(DataPoint):
-    text: str
-    chunk_size: int
-    chunk_index: int
-    cut_type: str
-    is_part_of: Document
-    contains: List[Entity] = None
+@pytest.fixture
+def mock_vector_engine():
+    """Create a mock vector engine."""
+    engine = AsyncMock()
+    engine.search = AsyncMock()
+    return engine
 
-    metadata: dict = {"index_fields": ["text"]}
 
+@pytest.mark.asyncio
+async def test_get_context_success(mock_vector_engine):
+    """Test successful retrieval of context."""
+    mock_result1 = MagicMock()
+    mock_result1.payload = {"text": "Steve Rodger"}
+    mock_result2 = MagicMock()
+    mock_result2.payload = {"text": "Mike Broski"}
 
-class TestRAGCompletionRetriever:
-    @pytest.mark.asyncio
-    async def test_rag_completion_context_simple(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_rag_completion_context_simple"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_rag_completion_context_simple"
-        )
-        cognee.config.data_root_directory(data_directory_path)
+    mock_vector_engine.search.return_value = [mock_result1, mock_result2]
 
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-        await setup()
+    retriever = CompletionRetriever(top_k=2)
 
-        document = TextDocument(
-            name="Steve Rodger's career",
-            raw_data_location="somewhere",
-            external_metadata="",
-            mime_type="text/plain",
-        )
+    with patch(
+        "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
 
-        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=[],
-        )
+    assert context == "Steve Rodger\nMike Broski"
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=2, include_payload=True
+    )
 
-        entities = [chunk1, chunk2, chunk3]
 
-        await add_data_points(entities)
+@pytest.mark.asyncio
+async def test_get_context_collection_not_found_error(mock_vector_engine):
+    """Test that CollectionNotFoundError is converted to NoDataError."""
+    mock_vector_engine.search.side_effect = CollectionNotFoundError("Collection not found")
 
-        retriever = CompletionRetriever()
+    retriever = CompletionRetriever()
 
-        context = await retriever.get_context("Mike")
+    with patch(
+        "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        with pytest.raises(NoDataError, match="No data found"):
+            await retriever.get_retrieved_objects("test query")
 
-        assert context == "Mike Broski", "Failed to get Mike Broski"
 
-    @pytest.mark.asyncio
-    async def test_rag_completion_context_complex(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_rag_completion_context_complex"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_rag_completion_context_complex"
-        )
-        cognee.config.data_root_directory(data_directory_path)
+@pytest.mark.asyncio
+async def test_get_context_empty_results(mock_vector_engine):
+    """Test that empty string is returned when no chunks are found."""
+    mock_vector_engine.search.return_value = []
 
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-        await setup()
+    retriever = CompletionRetriever()
 
-        document1 = TextDocument(
-            name="Employee List",
-            raw_data_location="somewhere",
-            external_metadata="",
-            mime_type="text/plain",
-        )
+    with patch(
+        "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        context = await retriever.get_context_from_objects("test query", [])
 
-        document2 = TextDocument(
-            name="Car List",
-            raw_data_location="somewhere",
-            external_metadata="",
-            mime_type="text/plain",
-        )
+    assert context == ""
 
-        chunk1 = DocumentChunk(
-            text="Steve Rodger",
-            chunk_size=2,
-            chunk_index=0,
-            cut_type="sentence_end",
-            is_part_of=document1,
-            contains=[],
-        )
-        chunk2 = DocumentChunk(
-            text="Mike Broski",
-            chunk_size=2,
-            chunk_index=1,
-            cut_type="sentence_end",
-            is_part_of=document1,
-            contains=[],
-        )
-        chunk3 = DocumentChunk(
-            text="Christina Mayer",
-            chunk_size=2,
-            chunk_index=2,
-            cut_type="sentence_end",
-            is_part_of=document1,
-            contains=[],
-        )
 
-        chunk4 = DocumentChunk(
-            text="Range Rover",
-            chunk_size=2,
-            chunk_index=0,
-            cut_type="sentence_end",
-            is_part_of=document2,
-            contains=[],
-        )
-        chunk5 = DocumentChunk(
-            text="Hyundai",
-            chunk_size=2,
-            chunk_index=1,
-            cut_type="sentence_end",
-            is_part_of=document2,
-            contains=[],
-        )
-        chunk6 = DocumentChunk(
-            text="Chrysler",
-            chunk_size=2,
-            chunk_index=2,
-            cut_type="sentence_end",
-            is_part_of=document2,
-            contains=[],
-        )
+@pytest.mark.asyncio
+async def test_get_context_top_k_limit(mock_vector_engine):
+    """Test that top_k parameter limits the number of results."""
+    mock_results = [MagicMock() for _ in range(2)]
+    for i, result in enumerate(mock_results):
+        result.payload = {"text": f"Chunk {i}"}
 
-        entities = [chunk1, chunk2, chunk3, chunk4, chunk5, chunk6]
+    mock_vector_engine.search.return_value = mock_results
 
-        await add_data_points(entities)
+    retriever = CompletionRetriever(top_k=2)
 
-        # TODO: top_k doesn't affect the output, it should be fixed.
-        retriever = CompletionRetriever(top_k=20)
+    with patch(
+        "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
 
-        context = await retriever.get_context("Christina")
+    assert context == "Chunk 0\nChunk 1"
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=2, include_payload=True
+    )
 
-        assert context[0:15] == "Christina Mayer", "Failed to get Christina Mayer"
 
-    @pytest.mark.asyncio
-    async def test_get_rag_completion_context_on_empty_graph(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".cognee_system/test_get_rag_completion_context_on_empty_graph",
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent,
-            ".data_storage/test_get_rag_completion_context_on_empty_graph",
+@pytest.mark.asyncio
+async def test_get_context_single_chunk(mock_vector_engine):
+    """Test get_context with single chunk result."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Single chunk text"}
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = CompletionRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
+
+    assert context == "Single chunk text"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_without_session(mock_vector_engine):
+    """Test get_completion without session caching."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Chunk text"}
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = CompletionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch("cognee.modules.retrieval.completion_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_from_context("test query", None, "")
+
+    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_vector_engine):
+    """Test get_completion with provided context."""
+    retriever = CompletionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.completion_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch("cognee.modules.retrieval.completion_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_from_context(
+            "test query", None, context="Provided context"
         )
-        cognee.config.data_root_directory(data_directory_path)
 
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert completion[0] == "Generated answer"
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_session(mock_vector_engine):
+    """Test get_completion with session caching enabled."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Chunk text"}
+    mock_vector_engine.search.return_value = [mock_result]
 
-        retriever = CompletionRetriever()
+    retriever = CompletionRetriever(session_id="test_session")
 
-        with pytest.raises(NoDataError):
-            await retriever.get_context("Christina Mayer")
+    mock_user = MagicMock()
+    mock_user.id = "test-user-id"
 
-        vector_engine = get_vector_engine()
-        await vector_engine.create_collection(
-            "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
+    with (
+        patch(
+            "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.get_conversation_history",
+            return_value="Previous conversation",
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.summarize_text",
+            return_value="Context summary",
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.save_conversation_history",
+        ) as mock_save,
+        patch("cognee.modules.retrieval.completion_retriever.CacheConfig") as mock_cache_config,
+        patch("cognee.modules.retrieval.completion_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_from_context(
+            "test query", [mock_result], "test"
         )
 
-        context = await retriever.get_context("Christina Mayer")
-        assert context == "", "Returned context should be empty on an empty graph"
+    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_session_no_user_id(mock_vector_engine):
+    """Test get_completion with session config but no user ID."""
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Chunk text"}
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = CompletionRetriever()
+
+    with (
+        patch(
+            "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.generate_completion",
+            return_value="Generated answer",
+        ),
+        patch("cognee.modules.retrieval.completion_retriever.CacheConfig") as mock_cache_config,
+        patch("cognee.modules.retrieval.completion_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_from_context("test query", None, "")
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+
+
+@pytest.mark.asyncio
+async def test_get_completion_with_response_model(mock_vector_engine):
+    """Test get_completion with custom response model."""
+    from pydantic import BaseModel
+
+    class TestModel(BaseModel):
+        answer: str
+
+    mock_result = MagicMock()
+    mock_result.payload = {"text": "Chunk text"}
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = CompletionRetriever(response_model=TestModel)
+
+    with (
+        patch(
+            "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+            return_value=mock_vector_engine,
+        ),
+        patch(
+            "cognee.modules.retrieval.completion_retriever.generate_completion",
+            return_value=TestModel(answer="Test answer"),
+        ),
+        patch("cognee.modules.retrieval.completion_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_from_context("test query", None, None)
+
+    assert isinstance(completion, list)
+    assert len(completion) == 1
+    assert isinstance(completion[0], TestModel)
+
+
+@pytest.mark.asyncio
+async def test_init_defaults():
+    """Test CompletionRetriever initialization with defaults."""
+    retriever = CompletionRetriever()
+
+    assert retriever.user_prompt_path == "context_for_question.txt"
+    assert retriever.system_prompt_path == "answer_simple_question.txt"
+    assert retriever.top_k == 1
+    assert retriever.system_prompt is None
+
+
+@pytest.mark.asyncio
+async def test_init_custom_params():
+    """Test CompletionRetriever initialization with custom parameters."""
+    retriever = CompletionRetriever(
+        user_prompt_path="custom_user.txt",
+        system_prompt_path="custom_system.txt",
+        system_prompt="Custom prompt",
+        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.top_k == 10
+
+
+@pytest.mark.asyncio
+async def test_get_context_missing_text_key(mock_vector_engine):
+    """Test get_context handles missing text key in payload."""
+    mock_result = MagicMock()
+    mock_result.payload = {"other_key": "value"}
+
+    mock_vector_engine.search.return_value = [mock_result]
+
+    retriever = CompletionRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.completion_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        with pytest.raises(KeyError):
+            objects = await retriever.get_retrieved_objects("test query")
+            context = await retriever.get_context_from_objects("test query", objects)
+            await retriever.get_completion_from_context("test query", objects, context)