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

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

@@ -1,201 +1,155 @@
-import os
 import pytest
-import pathlib
-from typing import List
-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 types import SimpleNamespace
+from unittest.mock import AsyncMock, patch, MagicMock
+
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
-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 chunk context."""
+    mock_result1 = MagicMock()
+    mock_result1.payload = {"text": "Steve Rodger", "chunk_index": 0}
+    mock_result2 = MagicMock()
+    mock_result2.payload = {"text": "Mike Broski", "chunk_index": 1}
 
-class TestChunksRetriever:
-    @pytest.mark.asyncio
-    async def test_chunk_context_simple(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_chunk_context_simple"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_chunk_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 = ChunksRetriever(top_k=5)
 
-        document = TextDocument(
-            name="Steve Rodger's career",
-            raw_data_location="somewhere",
-            external_metadata="",
-            mime_type="text/plain",
-        )
+    with patch(
+        "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
 
-        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 len(objects) == 2
+    assert objects[0].payload["text"] == "Steve Rodger"
+    assert objects[1].payload["text"] == "Mike Broski"
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=5, 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 = ChunksRetriever()
+    retriever = ChunksRetriever()
 
-        context = await retriever.get_context("Mike")
+    with patch(
+        "cognee.modules.retrieval.chunks_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[0]["text"] == "Mike Broski", "Failed to get Mike Broski"
 
-    @pytest.mark.asyncio
-    async def test_chunk_context_complex(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_chunk_context_complex"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_chunk_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 list 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 = ChunksRetriever()
 
-        document1 = TextDocument(
-            name="Employee List",
-            raw_data_location="somewhere",
-            external_metadata="",
-            mime_type="text/plain",
-        )
+    with patch(
+        "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
 
-        document2 = TextDocument(
-            name="Car List",
-            raw_data_location="somewhere",
-            external_metadata="",
-            mime_type="text/plain",
-        )
+    assert objects == []
 
-        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(3)]
+    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 = ChunksRetriever(top_k=3)
 
-        retriever = ChunksRetriever(top_k=20)
+    with patch(
+        "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        objects = await retriever.get_retrieved_objects("test query")
 
-        context = await retriever.get_context("Christina")
+    assert len(objects) == 3
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=3, include_payload=True
+    )
 
-        assert context[0]["text"] == "Christina Mayer", "Failed to get Christina Mayer"
 
-    @pytest.mark.asyncio
-    async def test_chunk_context_on_empty_graph(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_chunk_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_chunk_context_on_empty_graph"
-        )
-        cognee.config.data_root_directory(data_directory_path)
+@pytest.mark.asyncio
+async def test_get_context(mock_vector_engine):
+    """Test get_completion returns provided context."""
+    retriever = ChunksRetriever()
+
+    retrieved_objects = [
+        {"payload": {"text": "Steve Rodger"}},
+        {"payload": {"text": "Mike Broski"}},
+    ]
+    # Wrap the outer dictionary so payload is an attribute
+    mock_objects = [SimpleNamespace(**obj) for obj in retrieved_objects]
+
+    context = await retriever.get_context_from_objects("test query", retrieved_objects=mock_objects)
+
+    assert context == "Steve Rodger\nMike Broski"
+
+
+@pytest.mark.asyncio
+async def test_init_defaults():
+    """Test ChunksRetriever initialization with defaults."""
+    retriever = ChunksRetriever()
+
+    assert retriever.top_k == 5
+
+
+@pytest.mark.asyncio
+async def test_init_custom_top_k():
+    """Test ChunksRetriever initialization with custom top_k."""
+    retriever = ChunksRetriever(top_k=10)
+
+    assert retriever.top_k == 10
+
+
+@pytest.mark.asyncio
+async def test_init_none_top_k():
+    """Test ChunksRetriever initialization with None top_k."""
+    retriever = ChunksRetriever(top_k=None)
+
+    assert retriever.top_k is None
 
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
 
-        retriever = ChunksRetriever()
+@pytest.mark.asyncio
+async def test_get_context_empty_payload(mock_vector_engine):
+    """Test get_context handles empty payload."""
+    mock_vector_engine.search.return_value = []
 
-        with pytest.raises(NoDataError):
-            await retriever.get_context("Christina Mayer")
+    retriever = ChunksRetriever()
 
-        vector_engine = get_vector_engine()
-        await vector_engine.create_collection(
-            "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
+    with patch(
+        "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        retrieved_objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects(
+            "test query", retrieved_objects=retrieved_objects
         )
 
-        context = await retriever.get_context("Christina Mayer")
-        assert len(context) == 0, "Found chunks when none should exist"
+    assert context == ""

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

@@ -152,3 +152,341 @@ class TestConversationHistoryUtils:
                 assert result is True
                 call_kwargs = mock_cache.add_qa.call_args.kwargs
                 assert call_kwargs["session_id"] == "default_session"
+
+    @pytest.mark.asyncio
+    async def test_save_conversation_history_no_user_id(self):
+        """Test save_conversation_history returns False when user_id is None."""
+        session_user.set(None)
+
+        with patch("cognee.modules.retrieval.utils.session_cache.CacheConfig") as MockCacheConfig:
+            mock_config = MagicMock()
+            mock_config.caching = True
+            MockCacheConfig.return_value = mock_config
+
+            from cognee.modules.retrieval.utils.session_cache import (
+                save_conversation_history,
+            )
+
+            result = await save_conversation_history(
+                query="Test question",
+                context_summary="Test context",
+                answer="Test answer",
+            )
+
+            assert result is False
+
+    @pytest.mark.asyncio
+    async def test_save_conversation_history_caching_disabled(self):
+        """Test save_conversation_history returns False when caching is disabled."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        with patch("cognee.modules.retrieval.utils.session_cache.CacheConfig") as MockCacheConfig:
+            mock_config = MagicMock()
+            mock_config.caching = False
+            MockCacheConfig.return_value = mock_config
+
+            from cognee.modules.retrieval.utils.session_cache import (
+                save_conversation_history,
+            )
+
+            result = await save_conversation_history(
+                query="Test question",
+                context_summary="Test context",
+                answer="Test answer",
+            )
+
+            assert result is False
+
+    @pytest.mark.asyncio
+    async def test_save_conversation_history_cache_engine_none(self):
+        """Test save_conversation_history returns False when cache_engine is None."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=None):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    save_conversation_history,
+                )
+
+                result = await save_conversation_history(
+                    query="Test question",
+                    context_summary="Test context",
+                    answer="Test answer",
+                )
+
+                assert result is False
+
+    @pytest.mark.asyncio
+    async def test_save_conversation_history_cache_connection_error(self):
+        """Test save_conversation_history handles CacheConnectionError gracefully."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        from cognee.infrastructure.databases.exceptions import CacheConnectionError
+
+        mock_cache = create_mock_cache_engine([])
+        mock_cache.add_qa = AsyncMock(side_effect=CacheConnectionError("Connection failed"))
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=mock_cache):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    save_conversation_history,
+                )
+
+                result = await save_conversation_history(
+                    query="Test question",
+                    context_summary="Test context",
+                    answer="Test answer",
+                )
+
+                assert result is False
+
+    @pytest.mark.asyncio
+    async def test_save_conversation_history_generic_exception(self):
+        """Test save_conversation_history handles generic exceptions gracefully."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        mock_cache = create_mock_cache_engine([])
+        mock_cache.add_qa = AsyncMock(side_effect=ValueError("Unexpected error"))
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=mock_cache):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    save_conversation_history,
+                )
+
+                result = await save_conversation_history(
+                    query="Test question",
+                    context_summary="Test context",
+                    answer="Test answer",
+                )
+
+                assert result is False
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_no_user_id(self):
+        """Test get_conversation_history returns empty string when user_id is None."""
+        session_user.set(None)
+
+        with patch("cognee.modules.retrieval.utils.session_cache.CacheConfig") as MockCacheConfig:
+            mock_config = MagicMock()
+            mock_config.caching = True
+            MockCacheConfig.return_value = mock_config
+
+            from cognee.modules.retrieval.utils.session_cache import (
+                get_conversation_history,
+            )
+
+            result = await get_conversation_history(session_id="test_session")
+
+            assert result == ""
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_caching_disabled(self):
+        """Test get_conversation_history returns empty string when caching is disabled."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        with patch("cognee.modules.retrieval.utils.session_cache.CacheConfig") as MockCacheConfig:
+            mock_config = MagicMock()
+            mock_config.caching = False
+            MockCacheConfig.return_value = mock_config
+
+            from cognee.modules.retrieval.utils.session_cache import (
+                get_conversation_history,
+            )
+
+            result = await get_conversation_history(session_id="test_session")
+
+            assert result == ""
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_default_session(self):
+        """Test get_conversation_history uses 'default_session' when session_id is None."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        mock_cache = create_mock_cache_engine([])
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=mock_cache):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    get_conversation_history,
+                )
+
+                await get_conversation_history(session_id=None)
+
+                mock_cache.get_latest_qa.assert_called_once_with(str(user.id), "default_session")
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_cache_engine_none(self):
+        """Test get_conversation_history returns empty string when cache_engine is None."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=None):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    get_conversation_history,
+                )
+
+                result = await get_conversation_history(session_id="test_session")
+
+                assert result == ""
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_cache_connection_error(self):
+        """Test get_conversation_history handles CacheConnectionError gracefully."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        from cognee.infrastructure.databases.exceptions import CacheConnectionError
+
+        mock_cache = create_mock_cache_engine([])
+        mock_cache.get_latest_qa = AsyncMock(side_effect=CacheConnectionError("Connection failed"))
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=mock_cache):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    get_conversation_history,
+                )
+
+                result = await get_conversation_history(session_id="test_session")
+
+                assert result == ""
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_generic_exception(self):
+        """Test get_conversation_history handles generic exceptions gracefully."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        mock_cache = create_mock_cache_engine([])
+        mock_cache.get_latest_qa = AsyncMock(side_effect=ValueError("Unexpected error"))
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=mock_cache):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    get_conversation_history,
+                )
+
+                result = await get_conversation_history(session_id="test_session")
+
+                assert result == ""
+
+    @pytest.mark.asyncio
+    async def test_get_conversation_history_missing_keys(self):
+        """Test get_conversation_history handles missing keys in history entries."""
+        user = create_mock_user()
+        session_user.set(user)
+
+        mock_history = [
+            {
+                "time": "2024-01-15 10:30:45",
+                "question": "What is AI?",
+            },
+            {
+                "context": "AI is artificial intelligence",
+                "answer": "AI stands for Artificial Intelligence",
+            },
+            {},
+        ]
+        mock_cache = create_mock_cache_engine(mock_history)
+
+        cache_module = importlib.import_module(
+            "cognee.infrastructure.databases.cache.get_cache_engine"
+        )
+
+        with patch.object(cache_module, "get_cache_engine", return_value=mock_cache):
+            with patch(
+                "cognee.modules.retrieval.utils.session_cache.CacheConfig"
+            ) as MockCacheConfig:
+                mock_config = MagicMock()
+                mock_config.caching = True
+                MockCacheConfig.return_value = mock_config
+
+                from cognee.modules.retrieval.utils.session_cache import (
+                    get_conversation_history,
+                )
+
+                result = await get_conversation_history(session_id="test_session")
+
+                assert "Previous conversation:" in result
+                assert "[2024-01-15 10:30:45]" in result
+                assert "QUESTION: What is AI?" in result
+                assert "Unknown time" in result
+                assert "CONTEXT: AI is artificial intelligence" in result
+                assert "ANSWER: AI stands for Artificial Intelligence" in result