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

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

@@ -66,6 +66,12 @@ async def test_get_context_success(mock_edge):
     mock_graph_engine = AsyncMock()
     mock_graph_engine.is_empty = AsyncMock(return_value=False)
 
+    mock_node1 = MagicMock()
+    mock_node2 = MagicMock()
+    mock_edge.node1 = mock_node1
+    mock_edge.node2 = mock_node2
+    mock_edge.attributes = {"text": "mock edge"}
+
     with (
         patch(
             "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
@@ -76,11 +82,9 @@ async def test_get_context_success(mock_edge):
             return_value=[mock_edge],
         ),
     ):
-        context = await retriever.get_context("test query")
+        context = await retriever.get_context_from_objects("test query", [mock_edge])
 
-    assert isinstance(context, list)
-    assert len(context) == 1
-    assert context[0] == mock_edge
+    assert isinstance(context, str)
 
 
 @pytest.mark.asyncio
@@ -101,9 +105,9 @@ async def test_get_context_empty_results():
             return_value=[],
         ),
     ):
-        context = await retriever.get_context("test query")
+        context = await retriever.get_context_from_objects("test query", [])
 
-    assert context == []
+    assert context == ""
 
 
 @pytest.mark.asyncio
@@ -118,9 +122,9 @@ async def test_get_context_empty_graph():
         "cognee.modules.retrieval.graph_completion_retriever.get_graph_engine",
         return_value=mock_graph_engine,
     ):
-        context = await retriever.get_context("test query")
+        context = await retriever.get_context_from_objects("test query", [])
 
-    assert context == []
+    assert context == ""
 
 
 @pytest.mark.asyncio
@@ -184,21 +188,6 @@ async def test_init_none_top_k():
     assert retriever.top_k == 5  # None defaults to 5
 
 
-@pytest.mark.asyncio
-async def test_convert_retrieved_objects_to_context(mock_edge):
-    """Test convert_retrieved_objects_to_context method."""
-    retriever = GraphCompletionRetriever()
-
-    with patch(
-        "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
-        return_value="Resolved text",
-    ) as mock_resolve:
-        result = await retriever.convert_retrieved_objects_to_context([mock_edge])
-
-    assert result == "Resolved text"
-    mock_resolve.assert_awaited_once_with([mock_edge])
-
-
 @pytest.mark.asyncio
 async def test_get_completion_without_context(mock_edge):
     """Test get_completion retrieves context when not provided."""
@@ -232,7 +221,7 @@ async def test_get_completion_without_context(mock_edge):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query")
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -261,7 +250,9 @@ async def test_get_completion_with_provided_context(mock_edge):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", context=[mock_edge])
+        completion = await retriever.get_completion_from_context(
+            "test query", None, context="mock edge"
+        )
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -274,7 +265,7 @@ async def test_get_completion_with_session(mock_edge):
     mock_graph_engine = AsyncMock()
     mock_graph_engine.is_empty = AsyncMock(return_value=False)
 
-    retriever = GraphCompletionRetriever()
+    retriever = GraphCompletionRetriever(session_id="test_session")
 
     mock_user = MagicMock()
     mock_user.id = "test-user-id"
@@ -319,7 +310,7 @@ async def test_get_completion_with_session(mock_edge):
         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")
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -338,7 +329,7 @@ async def test_get_completion_with_response_model(mock_edge):
     mock_graph_engine = AsyncMock()
     mock_graph_engine.is_empty = AsyncMock(return_value=False)
 
-    retriever = GraphCompletionRetriever()
+    retriever = GraphCompletionRetriever(response_model=TestModel)
 
     with (
         patch(
@@ -365,7 +356,7 @@ async def test_get_completion_with_response_model(mock_edge):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", response_model=TestModel)
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -405,7 +396,7 @@ async def test_get_completion_empty_context(mock_edge):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query")
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -546,7 +537,7 @@ async def test_get_completion_with_save_interaction_no_completion(mock_edge):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query")
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -586,7 +577,7 @@ async def test_get_completion_with_save_interaction_no_context(mock_edge):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", context=None)
+        completion = await retriever.get_completion_from_context("test query", None, context=None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -640,7 +631,10 @@ async def test_get_completion_with_save_interaction_all_conditions_met(mock_edge
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", context=[mock_edge])
+        objects = await retriever.get_retrieved_objects("test query")
+        completion = await retriever.get_completion_from_context(
+            "test query", objects, context="mock_edge"
+        )
 
     assert isinstance(completion, list)
     assert len(completion) == 1

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

@@ -30,10 +30,13 @@ async def test_get_context_success(mock_vector_engine):
         "cognee.modules.retrieval.completion_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
 
     assert context == "Steve Rodger\nMike Broski"
-    mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=2)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=2, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
@@ -48,7 +51,7 @@ async def test_get_context_collection_not_found_error(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(NoDataError, match="No data found"):
-            await retriever.get_context("test query")
+            await retriever.get_retrieved_objects("test query")
 
 
 @pytest.mark.asyncio
@@ -62,7 +65,7 @@ async def test_get_context_empty_results(mock_vector_engine):
         "cognee.modules.retrieval.completion_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        context = await retriever.get_context_from_objects("test query", [])
 
     assert context == ""
 
@@ -82,10 +85,13 @@ async def test_get_context_top_k_limit(mock_vector_engine):
         "cognee.modules.retrieval.completion_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
 
     assert context == "Chunk 0\nChunk 1"
-    mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=2)
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_text", "test query", limit=2, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
@@ -101,7 +107,8 @@ async def test_get_context_single_chunk(mock_vector_engine):
         "cognee.modules.retrieval.completion_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
 
     assert context == "Single chunk text"
 
@@ -130,7 +137,7 @@ async def test_get_completion_without_session(mock_vector_engine):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query")
+        completion = await retriever.get_completion_from_context("test query", None, "")
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -153,7 +160,9 @@ async def test_get_completion_with_provided_context(mock_vector_engine):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", context="Provided context")
+        completion = await retriever.get_completion_from_context(
+            "test query", None, context="Provided context"
+        )
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -167,7 +176,7 @@ async def test_get_completion_with_session(mock_vector_engine):
     mock_result.payload = {"text": "Chunk text"}
     mock_vector_engine.search.return_value = [mock_result]
 
-    retriever = CompletionRetriever()
+    retriever = CompletionRetriever(session_id="test_session")
 
     mock_user = MagicMock()
     mock_user.id = "test-user-id"
@@ -200,7 +209,9 @@ async def test_get_completion_with_session(mock_vector_engine):
         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")
+        completion = await retriever.get_completion_from_context(
+            "test query", [mock_result], "test"
+        )
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -234,7 +245,7 @@ async def test_get_completion_with_session_no_user_id(mock_vector_engine):
         mock_cache_config.return_value = mock_config
         mock_session_user.get.return_value = None  # No user
 
-        completion = await retriever.get_completion("test query")
+        completion = await retriever.get_completion_from_context("test query", None, "")
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -252,7 +263,7 @@ async def test_get_completion_with_response_model(mock_vector_engine):
     mock_result.payload = {"text": "Chunk text"}
     mock_vector_engine.search.return_value = [mock_result]
 
-    retriever = CompletionRetriever()
+    retriever = CompletionRetriever(response_model=TestModel)
 
     with (
         patch(
@@ -269,7 +280,7 @@ async def test_get_completion_with_response_model(mock_vector_engine):
         mock_config.caching = False
         mock_cache_config.return_value = mock_config
 
-        completion = await retriever.get_completion("test query", response_model=TestModel)
+        completion = await retriever.get_completion_from_context("test query", None, None)
 
     assert isinstance(completion, list)
     assert len(completion) == 1
@@ -318,4 +329,6 @@ async def test_get_context_missing_text_key(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(KeyError):
-            await retriever.get_context("test query")
+            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)

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

@@ -1,5 +1,6 @@
 import pytest
 from unittest.mock import AsyncMock, patch, MagicMock
+from types import SimpleNamespace
 
 from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
@@ -30,12 +31,16 @@ async def test_get_context_success(mock_vector_engine):
         "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
 
-    assert len(context) == 2
-    assert context[0]["text"] == "S.R."
-    assert context[1]["text"] == "M.B."
-    mock_vector_engine.search.assert_awaited_once_with("TextSummary_text", "test query", limit=5)
+    assert len(completion) == 2
+    assert completion[0]["text"] == "S.R."
+    assert completion[1]["text"] == "M.B."
+    mock_vector_engine.search.assert_awaited_once_with(
+        "TextSummary_text", "test query", limit=5, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
@@ -50,11 +55,11 @@ async def test_get_context_collection_not_found_error(mock_vector_engine):
         return_value=mock_vector_engine,
     ):
         with pytest.raises(NoDataError, match="No data found"):
-            await retriever.get_context("test query")
+            await retriever.get_retrieved_objects("test query")
 
 
 @pytest.mark.asyncio
-async def test_get_context_empty_results(mock_vector_engine):
+async def test_get_objects_empty_results(mock_vector_engine):
     """Test that empty list is returned when no summaries are found."""
     mock_vector_engine.search.return_value = []
 
@@ -64,13 +69,13 @@ async def test_get_context_empty_results(mock_vector_engine):
         "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
 
-    assert context == []
+    assert objects == []
 
 
 @pytest.mark.asyncio
-async def test_get_context_top_k_limit(mock_vector_engine):
+async def test_get_objects_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):
@@ -84,21 +89,25 @@ async def test_get_context_top_k_limit(mock_vector_engine):
         "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
 
-    assert len(context) == 3
-    mock_vector_engine.search.assert_awaited_once_with("TextSummary_text", "test query", limit=3)
+    assert len(objects) == 3
+    mock_vector_engine.search.assert_awaited_once_with(
+        "TextSummary_text", "test query", limit=3, include_payload=True
+    )
 
 
 @pytest.mark.asyncio
-async def test_get_completion_with_context(mock_vector_engine):
+async def test_get_context_with_objects(mock_vector_engine):
     """Test get_completion returns provided context."""
     retriever = SummariesRetriever()
 
-    provided_context = [{"text": "S.R."}, {"text": "M.B."}]
-    completion = await retriever.get_completion("test query", context=provided_context)
+    provided_context = {"text": "S.R."}
+    sn = SimpleNamespace()
+    sn.payload = provided_context
+    completion = await retriever.get_context_from_objects("test query", retrieved_objects=[sn])
 
-    assert completion == provided_context
+    assert completion == provided_context["text"]
 
 
 @pytest.mark.asyncio
@@ -114,7 +123,9 @@ async def test_get_completion_without_context(mock_vector_engine):
         "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        completion = await retriever.get_completion("test query")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
 
     assert len(completion) == 1
     assert completion[0]["text"] == "S.R."
@@ -137,7 +148,7 @@ async def test_init_custom_top_k():
 
 
 @pytest.mark.asyncio
-async def test_get_context_empty_payload(mock_vector_engine):
+async def test_get_objects_empty_payload(mock_vector_engine):
     """Test get_context handles empty payload."""
     mock_result = MagicMock()
     mock_result.payload = {}
@@ -150,10 +161,10 @@ async def test_get_context_empty_payload(mock_vector_engine):
         "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        context = await retriever.get_context("test query")
+        objects = await retriever.get_retrieved_objects("test query")
 
-    assert len(context) == 1
-    assert context[0] == {}
+    assert len(objects) == 1
+    assert objects[0].payload == {}
 
 
 @pytest.mark.asyncio
@@ -163,31 +174,15 @@ async def test_get_completion_with_session_id(mock_vector_engine):
     mock_result.payload = {"text": "S.R."}
     mock_vector_engine.search.return_value = [mock_result]
 
-    retriever = SummariesRetriever()
+    retriever = SummariesRetriever(session_id="test_session")
 
     with patch(
         "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        completion = await retriever.get_completion("test query", session_id="test_session")
+        objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects("test query", objects)
+        completion = await retriever.get_completion_from_context("test query", objects, context)
 
     assert len(completion) == 1
     assert completion[0]["text"] == "S.R."
-
-
-@pytest.mark.asyncio
-async def test_get_completion_with_kwargs(mock_vector_engine):
-    """Test get_completion accepts additional kwargs."""
-    mock_result = MagicMock()
-    mock_result.payload = {"text": "S.R."}
-    mock_vector_engine.search.return_value = [mock_result]
-
-    retriever = SummariesRetriever()
-
-    with patch(
-        "cognee.modules.retrieval.summaries_retriever.get_vector_engine",
-        return_value=mock_vector_engine,
-    ):
-        completion = await retriever.get_completion("test query", extra_param="value")
-
-    assert len(completion) == 1