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

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

@@ -9,7 +9,7 @@ def test_node_initialization():
     """Test that a Node is initialized correctly."""
     node = Node("node1", {"attr1": "value1"}, dimension=2)
     assert node.id == "node1"
-    assert node.attributes == {"attr1": "value1", "vector_distance": 3.5}
+    assert node.attributes == {"attr1": "value1", "vector_distance": None}
     assert len(node.status) == 2
     assert np.all(node.status == 1)
 
@@ -86,6 +86,46 @@ def test_node_hash():
     assert hash(node) == hash("node1")
 
 
+def test_node_vector_distance_stays_none():
+    """Test that vector_distance remains None when no distances are passed."""
+    node = Node("node1")
+    assert node.attributes.get("vector_distance") is None
+
+    # Verify it stays None even after other operations
+    node.add_attribute("other_attr", "value")
+    assert node.attributes.get("vector_distance") is None
+
+
+def test_node_vector_distance_with_custom_attributes():
+    """Test that vector_distance is None even when node has custom attributes."""
+    node = Node("node1", {"custom": "value", "another": 42})
+    assert node.attributes.get("vector_distance") is None
+    assert node.attributes["custom"] == "value"
+    assert node.attributes["another"] == 42
+
+
+def test_edge_vector_distance_stays_none():
+    """Test that vector_distance remains None when no distances are passed."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2)
+    assert edge.attributes.get("vector_distance") is None
+
+    # Verify it stays None even after other operations
+    edge.add_attribute("other_attr", "value")
+    assert edge.attributes.get("vector_distance") is None
+
+
+def test_edge_vector_distance_with_custom_attributes():
+    """Test that vector_distance is None even when edge has custom attributes."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, {"weight": 5, "type": "test"})
+    assert edge.attributes.get("vector_distance") is None
+    assert edge.attributes["weight"] == 5
+    assert edge.attributes["type"] == "test"
+
+
 ### Tests for Edge ###
 
 
@@ -96,7 +136,7 @@ def test_edge_initialization():
     edge = Edge(node1, node2, {"weight": 10}, directed=False, dimension=2)
     assert edge.node1 == node1
     assert edge.node2 == node2
-    assert edge.attributes == {"vector_distance": 3.5, "weight": 10}
+    assert edge.attributes == {"vector_distance": None, "weight": 10}
     assert edge.directed is False
     assert len(edge.status) == 2
     assert np.all(edge.status == 1)

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 == []