Merge dicts directly

topoteretes · Nov 19, 2024 · b18f748 · b18f748
1 parent fde56f0
commit b18f748
Showing 1 changed file with 5 additions and 20 deletions.
diff --git a/cognee/modules/graph/utils/get_model_instance_from_graph.py b/cognee/modules/graph/utils/get_model_instance_from_graph.py
@@ -6,20 +6,6 @@
 from cognee.modules.storage.utils import copy_model
 
 
-def merge_dicts(dict1: dict, dict2: dict, agg_fn: Callable) -> dict:
-    merged_dict = {}
-    for key, value in dict1.items():
-        if key in dict2:
-            merged_dict[key] = agg_fn(value, dict2[key])
-        else:
-            merged_dict[key] = value
-
-    for key, value in dict2.items():
-        if key not in merged_dict:
-            merged_dict[key] = value
-    return merged_dict
-
-
 def get_model_instance_from_graph(
     nodes: list[DataPoint],
     edges: list[tuple[str, str, str, dict[str, str]]],
@@ -38,12 +24,11 @@ def get_model_instance_from_graph(
                 type(source_node),
                 {edge_label: (list[type(target_node)], PydanticUndefined)},
             )
-            new_model_dict = merge_dicts(
-                source_node.model_dump(),
-                {edge_label: [target_node]},
-                lambda a, b: a + b,
-            )
-            node_map[source_node_id] = NewModel(**new_model_dict)
+            source_node_dict = source_node.model_dump()
+            source_node_edge_label_values = source_node_dict.get(edge_label, [])
+            source_node_dict[edge_label] = source_node_edge_label_values + [target_node]
+
+            node_map[source_node_id] = NewModel(**source_node_dict)
         else:
             NewModel = copy_model(
                 type(source_node), {edge_label: (type(target_node), PydanticUndefined)}