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

cognee/tests/unit/modules/graph/cognee_graph_test.py

@@ -1,6 +1,7 @@
 import pytest
 from unittest.mock import AsyncMock
 
+from cognee.modules.engine.utils.generate_edge_id import generate_edge_id
 from cognee.modules.graph.exceptions import EntityNotFoundError, EntityAlreadyExistsError
 from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node
@@ -200,6 +201,37 @@ async def test_project_graph_from_db_empty_graph(setup_graph, mock_adapter):
         )
 
 
+@pytest.mark.asyncio
+async def test_project_graph_from_db_stores_triplet_penalty_on_graph(mock_adapter):
+    """Test that project_graph_from_db stores triplet_distance_penalty on the graph."""
+    from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
+
+    nodes_data = [("1", {"name": "Node1"})]
+    edges_data = [("1", "1", "SELF", {})]
+
+    mock_adapter.get_graph_data = AsyncMock(return_value=(nodes_data, edges_data))
+
+    graph = CogneeGraph()
+    custom_penalty = 5.0
+    await graph.project_graph_from_db(
+        adapter=mock_adapter,
+        node_properties_to_project=["name"],
+        edge_properties_to_project=[],
+        triplet_distance_penalty=custom_penalty,
+    )
+
+    assert graph.triplet_distance_penalty == custom_penalty
+
+    graph2 = CogneeGraph()
+    await graph2.project_graph_from_db(
+        adapter=mock_adapter,
+        node_properties_to_project=["name"],
+        edge_properties_to_project=[],
+    )
+
+    assert graph2.triplet_distance_penalty == 3.5
+
+
 @pytest.mark.asyncio
 async def test_project_graph_from_db_missing_nodes(setup_graph, mock_adapter):
     """Test that edges referencing missing nodes raise error."""
@@ -241,8 +273,8 @@ async def test_map_vector_distances_to_graph_nodes(setup_graph):
 
     await graph.map_vector_distances_to_graph_nodes(node_distances)
 
-    assert graph.get_node("1").attributes.get("vector_distance") == 0.95
-    assert graph.get_node("2").attributes.get("vector_distance") == 0.87
+    assert graph.get_node("1").attributes.get("vector_distance") == [0.95]
+    assert graph.get_node("2").attributes.get("vector_distance") == [0.87]
 
 
 @pytest.mark.asyncio
@@ -266,9 +298,9 @@ async def test_map_vector_distances_partial_node_coverage(setup_graph):
 
     await graph.map_vector_distances_to_graph_nodes(node_distances)
 
-    assert graph.get_node("1").attributes.get("vector_distance") == 0.95
-    assert graph.get_node("2").attributes.get("vector_distance") == 0.87
-    assert graph.get_node("3").attributes.get("vector_distance") == 3.5
+    assert graph.get_node("1").attributes.get("vector_distance") == [0.95]
+    assert graph.get_node("2").attributes.get("vector_distance") == [0.87]
+    assert graph.get_node("3").attributes.get("vector_distance") == [3.5]
 
 
 @pytest.mark.asyncio
@@ -298,10 +330,36 @@ async def test_map_vector_distances_multiple_categories(setup_graph):
 
     await graph.map_vector_distances_to_graph_nodes(node_distances)
 
-    assert graph.get_node("1").attributes.get("vector_distance") == 0.95
-    assert graph.get_node("2").attributes.get("vector_distance") == 0.87
-    assert graph.get_node("3").attributes.get("vector_distance") == 0.92
-    assert graph.get_node("4").attributes.get("vector_distance") == 3.5
+    assert graph.get_node("1").attributes.get("vector_distance") == [0.95]
+    assert graph.get_node("2").attributes.get("vector_distance") == [0.87]
+    assert graph.get_node("3").attributes.get("vector_distance") == [0.92]
+    assert graph.get_node("4").attributes.get("vector_distance") == [3.5]
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_to_graph_nodes_multi_query(setup_graph):
+    """Test mapping vector distances with multiple queries."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    node_distances = {
+        "Entity_name": [
+            [MockScoredResult("1", 0.95)],  # query 0
+            [MockScoredResult("2", 0.87)],  # query 1
+        ]
+    }
+
+    await graph.map_vector_distances_to_graph_nodes(node_distances, query_list_length=2)
+
+    assert graph.get_node("1").attributes.get("vector_distance") == [0.95, 3.5]
+    assert graph.get_node("2").attributes.get("vector_distance") == [3.5, 0.87]
+    assert graph.get_node("3").attributes.get("vector_distance") == [3.5, 3.5]
 
 
 @pytest.mark.asyncio
@@ -322,12 +380,12 @@ async def test_map_vector_distances_to_graph_edges_with_payload(setup_graph):
     graph.add_edge(edge)
 
     edge_distances = [
-        MockScoredResult("e1", 0.92, payload={"text": "CONNECTS_TO"}),
+        MockScoredResult(generate_edge_id("CONNECTS_TO"), 0.92, payload={"text": "CONNECTS_TO"}),
     ]
 
     await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
 
-    assert graph.edges[0].attributes.get("vector_distance") == 0.92
+    assert graph.edges[0].attributes.get("vector_distance") == [0.92]
 
 
 @pytest.mark.asyncio
@@ -347,14 +405,15 @@ async def test_map_vector_distances_partial_edge_coverage(setup_graph):
     graph.add_edge(edge1)
     graph.add_edge(edge2)
 
+    edge_1_text = "CONNECTS_TO"
     edge_distances = [
-        MockScoredResult("e1", 0.92, payload={"text": "CONNECTS_TO"}),
+        MockScoredResult(generate_edge_id(edge_1_text), 0.92, payload={"text": edge_1_text}),
     ]
 
     await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
 
-    assert graph.edges[0].attributes.get("vector_distance") == 0.92
-    assert graph.edges[1].attributes.get("vector_distance") == 3.5
+    assert graph.edges[0].attributes.get("vector_distance") == [0.92]
+    assert graph.edges[1].attributes.get("vector_distance") == [3.5]
 
 
 @pytest.mark.asyncio
@@ -374,13 +433,14 @@ async def test_map_vector_distances_edges_fallback_to_relationship_type(setup_gr
     )
     graph.add_edge(edge)
 
+    edge_text = "KNOWS"
     edge_distances = [
-        MockScoredResult("e1", 0.85, payload={"text": "KNOWS"}),
+        MockScoredResult(generate_edge_id(edge_text), 0.85, payload={"text": edge_text}),
     ]
 
     await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
 
-    assert graph.edges[0].attributes.get("vector_distance") == 0.85
+    assert graph.edges[0].attributes.get("vector_distance") == [0.85]
 
 
 @pytest.mark.asyncio
@@ -400,18 +460,19 @@ async def test_map_vector_distances_no_edge_matches(setup_graph):
     )
     graph.add_edge(edge)
 
+    edge_text = "SOME_OTHER_EDGE"
     edge_distances = [
-        MockScoredResult("e1", 0.92, payload={"text": "SOME_OTHER_EDGE"}),
+        MockScoredResult(generate_edge_id(edge_text), 0.92, payload={"text": edge_text}),
     ]
 
     await graph.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
 
-    assert graph.edges[0].attributes.get("vector_distance") == 3.5
+    assert graph.edges[0].attributes.get("vector_distance") == [3.5]
 
 
 @pytest.mark.asyncio
 async def test_map_vector_distances_none_returns_early(setup_graph):
-    """Test that edge_distances=None returns early without error."""
+    """Test that edge_distances=None returns early without error and vector_distance is set to default penalty."""
     graph = setup_graph
     graph.add_node(Node("1"))
     graph.add_node(Node("2"))
@@ -419,7 +480,91 @@ async def test_map_vector_distances_none_returns_early(setup_graph):
 
     await graph.map_vector_distances_to_graph_edges(edge_distances=None)
 
-    assert graph.edges[0].attributes.get("vector_distance") == 3.5
+    assert graph.edges[0].attributes.get("vector_distance") == [3.5]
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_empty_nodes_returns_early(setup_graph):
+    """Test that node_distances={} returns early without error and vector_distance is set to default penalty."""
+    graph = setup_graph
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    await graph.map_vector_distances_to_graph_nodes({})
+
+    assert node1.attributes.get("vector_distance") == [3.5]
+    assert node2.attributes.get("vector_distance") == [3.5]
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_to_graph_edges_multi_query(setup_graph):
+    """Test mapping edge distances with multiple queries."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    edge1 = Edge(node1, node2, attributes={"edge_text": "A"})
+    edge2 = Edge(node2, node3, attributes={"edge_text": "B"})
+    graph.add_edge(edge1)
+    graph.add_edge(edge2)
+
+    edge_1_text = "A"
+    edge_2_text = "B"
+    edge_distances = [
+        [
+            MockScoredResult(generate_edge_id(edge_1_text), 0.1, payload={"text": edge_1_text})
+        ],  # query 0
+        [
+            MockScoredResult(generate_edge_id(edge_2_text), 0.2, payload={"text": edge_2_text})
+        ],  # query 1
+    ]
+
+    await graph.map_vector_distances_to_graph_edges(
+        edge_distances=edge_distances, query_list_length=2
+    )
+
+    assert graph.edges[0].attributes.get("vector_distance") == [0.1, 3.5]
+    assert graph.edges[1].attributes.get("vector_distance") == [3.5, 0.2]
+
+
+@pytest.mark.asyncio
+async def test_map_vector_distances_to_graph_edges_preserves_unmapped_indices(setup_graph):
+    """Test that unmapped indices in multi-query mode stay at default penalty."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    edge1 = Edge(node1, node2, attributes={"edge_text": "A"})
+    edge2 = Edge(node2, node3, attributes={"edge_text": "B"})
+    graph.add_edge(edge1)
+    graph.add_edge(edge2)
+
+    edge_1_text = "A"
+    edge_distances = [
+        [
+            MockScoredResult(generate_edge_id(edge_1_text), 0.1, payload={"text": edge_1_text})
+        ],  # query 0: only edge1 mapped
+        [],  # query 1: no edges mapped
+    ]
+
+    await graph.map_vector_distances_to_graph_edges(
+        edge_distances=edge_distances, query_list_length=2
+    )
+
+    assert graph.edges[0].attributes.get("vector_distance") == [0.1, 3.5]
+    assert graph.edges[1].attributes.get("vector_distance") == [3.5, 3.5]
 
 
 @pytest.mark.asyncio
@@ -432,10 +577,10 @@ async def test_calculate_top_triplet_importances(setup_graph):
     node3 = Node("3")
     node4 = Node("4")
 
-    node1.add_attribute("vector_distance", 0.9)
-    node2.add_attribute("vector_distance", 0.8)
-    node3.add_attribute("vector_distance", 0.7)
-    node4.add_attribute("vector_distance", 0.6)
+    node1.add_attribute("vector_distance", [0.9])
+    node2.add_attribute("vector_distance", [0.8])
+    node3.add_attribute("vector_distance", [0.7])
+    node4.add_attribute("vector_distance", [0.6])
 
     graph.add_node(node1)
     graph.add_node(node2)
@@ -446,9 +591,9 @@ async def test_calculate_top_triplet_importances(setup_graph):
     edge2 = Edge(node2, node3)
     edge3 = Edge(node3, node4)
 
-    edge1.add_attribute("vector_distance", 0.85)
-    edge2.add_attribute("vector_distance", 0.75)
-    edge3.add_attribute("vector_distance", 0.65)
+    edge1.add_attribute("vector_distance", [0.85])
+    edge2.add_attribute("vector_distance", [0.75])
+    edge3.add_attribute("vector_distance", [0.65])
 
     graph.add_edge(edge1)
     graph.add_edge(edge2)
@@ -464,7 +609,112 @@ async def test_calculate_top_triplet_importances(setup_graph):
 
 @pytest.mark.asyncio
 async def test_calculate_top_triplet_importances_default_distances(setup_graph):
-    """Test calculating importances when nodes/edges have no vector distances."""
+    """Test that vector_distance stays None when no distances are passed and calculate_top_triplet_importances handles it."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    edge = Edge(node1, node2)
+    graph.add_edge(edge)
+
+    # Verify vector_distance is None when no distances are passed
+    assert node1.attributes.get("vector_distance") is None
+    assert node2.attributes.get("vector_distance") is None
+    assert edge.attributes.get("vector_distance") is None
+
+    # When no distances are set, calculate_top_triplet_importances should handle None
+    # by either raising an error or skipping edges with None distances
+    with pytest.raises(ValueError):
+        await graph.calculate_top_triplet_importances(k=1)
+
+
+@pytest.mark.asyncio
+async def test_calculate_top_triplet_importances_single_query_via_helper(setup_graph):
+    """Test calculating top triplet importances for a single query index."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    node1.add_attribute("vector_distance", [0.1])
+    node2.add_attribute("vector_distance", [0.2])
+    node3.add_attribute("vector_distance", [0.3])
+
+    edge1 = Edge(node1, node2)
+    edge2 = Edge(node2, node3)
+    graph.add_edge(edge1)
+    graph.add_edge(edge2)
+
+    edge1.add_attribute("vector_distance", [0.3])
+    edge2.add_attribute("vector_distance", [0.4])
+
+    results = await graph.calculate_top_triplet_importances(k=1, query_list_length=1)
+    assert len(results) == 1
+    assert len(results[0]) == 1
+    assert results[0][0] == edge1
+
+
+@pytest.mark.asyncio
+async def test_calculate_top_triplet_importances_multi_query(setup_graph):
+    """Test calculating top triplet importances with multiple queries."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    graph.add_node(node1)
+    graph.add_node(node2)
+    graph.add_node(node3)
+
+    edge_a = Edge(node1, node2)
+    edge_b = Edge(node2, node3)
+    graph.add_edge(edge_a)
+    graph.add_edge(edge_b)
+
+    node1.add_attribute("vector_distance", [0.1, 0.9])
+    node2.add_attribute("vector_distance", [0.1, 0.9])
+    node3.add_attribute("vector_distance", [0.9, 0.1])
+    edge_a.add_attribute("vector_distance", [0.1, 0.9])
+    edge_b.add_attribute("vector_distance", [0.9, 0.1])
+
+    results = await graph.calculate_top_triplet_importances(k=1, query_list_length=2)
+
+    assert len(results) == 2
+    assert results[0][0] == edge_a
+    assert results[1][0] == edge_b
+
+
+@pytest.mark.asyncio
+async def test_calculate_top_triplet_importances_raises_on_short_list(setup_graph):
+    """Test that scoring raises ValueError when list is too short for query_index."""
+    graph = setup_graph
+
+    node1 = Node("1")
+    node2 = Node("2")
+    graph.add_node(node1)
+    graph.add_node(node2)
+
+    node1.add_attribute("vector_distance", [0.1])
+    node2.add_attribute("vector_distance", [0.2])
+
+    edge = Edge(node1, node2)
+    edge.add_attribute("vector_distance", [0.3])
+    graph.add_edge(edge)
+
+    with pytest.raises(ValueError):
+        await graph.calculate_top_triplet_importances(k=1, query_list_length=2)
+
+
+@pytest.mark.asyncio
+async def test_calculate_top_triplet_importances_raises_on_missing_attribute(setup_graph):
+    """Test that scoring raises error when vector_distance is missing."""
     graph = setup_graph
 
     node1 = Node("1")
@@ -472,10 +722,58 @@ async def test_calculate_top_triplet_importances_default_distances(setup_graph):
     graph.add_node(node1)
     graph.add_node(node2)
 
+    del node1.attributes["vector_distance"]
+    del node2.attributes["vector_distance"]
+
     edge = Edge(node1, node2)
+    del edge.attributes["vector_distance"]
     graph.add_edge(edge)
 
-    top_triplets = await graph.calculate_top_triplet_importances(k=1)
+    with pytest.raises(ValueError):
+        await graph.calculate_top_triplet_importances(k=1, query_list_length=1)
+
+
+def test_normalize_query_distance_lists_flat_list_single_query(setup_graph):
+    """Test that flat list is normalized to list-of-lists with length 1 for single-query mode."""
+    graph = setup_graph
+    flat_list = [MockScoredResult("node1", 0.95), MockScoredResult("node2", 0.87)]
+
+    result = graph._normalize_query_distance_lists(flat_list, query_list_length=None, name="test")
+
+    assert len(result) == 1
+    assert result[0] == flat_list
+
+
+def test_normalize_query_distance_lists_nested_list_batch_mode(setup_graph):
+    """Test that nested list is used as-is when query_list_length matches."""
+    graph = setup_graph
+    nested_list = [
+        [MockScoredResult("node1", 0.95)],
+        [MockScoredResult("node2", 0.87)],
+    ]
+
+    result = graph._normalize_query_distance_lists(nested_list, query_list_length=2, name="test")
+
+    assert len(result) == 2
+    assert result == nested_list
+
+
+def test_normalize_query_distance_lists_raises_on_length_mismatch(setup_graph):
+    """Test that ValueError is raised when nested list length doesn't match query_list_length."""
+    graph = setup_graph
+    nested_list = [
+        [MockScoredResult("node1", 0.95)],
+        [MockScoredResult("node2", 0.87)],
+    ]
+
+    with pytest.raises(ValueError, match="test has 2 query lists, but query_list_length is 3"):
+        graph._normalize_query_distance_lists(nested_list, query_list_length=3, name="test")
+
+
+def test_normalize_query_distance_lists_empty_list(setup_graph):
+    """Test that empty list returns empty list."""
+    graph = setup_graph
+
+    result = graph._normalize_query_distance_lists([], query_list_length=None, name="test")
 
-    assert len(top_triplets) == 1
-    assert top_triplets[0] == edge
+    assert result == []

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

@@ -1,4 +1,5 @@
 import pytest
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, patch, MagicMock
 
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
@@ -30,12 +31,14 @@ async def test_get_context_success(mock_vector_engine):
         "cognee.modules.retrieval.chunks_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) == 2
-    assert context[0]["text"] == "Steve Rodger"
-    assert context[1]["text"] == "Mike Broski"
-    mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=5)
+    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
+    )
 
 
 @pytest.mark.asyncio
@@ -50,7 +53,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
@@ -64,9 +67,9 @@ async def test_get_context_empty_results(mock_vector_engine):
         "cognee.modules.retrieval.chunks_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
@@ -84,40 +87,29 @@ async def test_get_context_top_k_limit(mock_vector_engine):
         "cognee.modules.retrieval.chunks_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("DocumentChunk_text", "test query", limit=3)
+    assert len(objects) == 3
+    mock_vector_engine.search.assert_awaited_once_with(
+        "DocumentChunk_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(mock_vector_engine):
     """Test get_completion returns provided context."""
     retriever = ChunksRetriever()
 
-    provided_context = [{"text": "Steve Rodger"}, {"text": "Mike Broski"}]
-    completion = await retriever.get_completion("test query", context=provided_context)
+    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]
 
-    assert completion == provided_context
+    context = await retriever.get_context_from_objects("test query", retrieved_objects=mock_objects)
 
-
-@pytest.mark.asyncio
-async def test_get_completion_without_context(mock_vector_engine):
-    """Test get_completion retrieves context when not provided."""
-    mock_result = MagicMock()
-    mock_result.payload = {"text": "Steve Rodger"}
-    mock_vector_engine.search.return_value = [mock_result]
-
-    retriever = ChunksRetriever()
-
-    with patch(
-        "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
-        return_value=mock_vector_engine,
-    ):
-        completion = await retriever.get_completion("test query")
-
-    assert len(completion) == 1
-    assert completion[0]["text"] == "Steve Rodger"
+    assert context == "Steve Rodger\nMike Broski"
 
 
 @pytest.mark.asyncio
@@ -147,29 +139,7 @@ async def test_init_none_top_k():
 @pytest.mark.asyncio
 async def test_get_context_empty_payload(mock_vector_engine):
     """Test get_context handles empty payload."""
-    mock_result = MagicMock()
-    mock_result.payload = {}
-
-    mock_vector_engine.search.return_value = [mock_result]
-
-    retriever = ChunksRetriever()
-
-    with patch(
-        "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
-        return_value=mock_vector_engine,
-    ):
-        context = await retriever.get_context("test query")
-
-    assert len(context) == 1
-    assert context[0] == {}
-
-
-@pytest.mark.asyncio
-async def test_get_completion_with_session_id(mock_vector_engine):
-    """Test get_completion with session_id parameter."""
-    mock_result = MagicMock()
-    mock_result.payload = {"text": "Steve Rodger"}
-    mock_vector_engine.search.return_value = [mock_result]
+    mock_vector_engine.search.return_value = []
 
     retriever = ChunksRetriever()
 
@@ -177,7 +147,9 @@ async def test_get_completion_with_session_id(mock_vector_engine):
         "cognee.modules.retrieval.chunks_retriever.get_vector_engine",
         return_value=mock_vector_engine,
     ):
-        completion = await retriever.get_completion("test query", session_id="test_session")
+        retrieved_objects = await retriever.get_retrieved_objects("test query")
+        context = await retriever.get_context_from_objects(
+            "test query", retrieved_objects=retrieved_objects
+        )
 
-    assert len(completion) == 1
-    assert completion[0]["text"] == "Steve Rodger"
+    assert context == ""