Changes to make OpenAI key work for dataStoreAgent

backend/src/utils/semantic_layer_builder.py

@@ -123,8 +123,11 @@ async def enriched_rel_properties(llm, model, finalised_graph_structure):
 
     # Merge properties
     for new_rel in enriched_rel_properties["relProperties"]:
-        relationship_type = new_rel["relType"]
-        properties_to_add = new_rel["property"]
+        #FS-67 - This change was done as relType is unavailable here for mistral.
+        relationship_type = new_rel.get("relType", None)
+        if relationship_type is None:
+            continue  # Skip if there's no relationship type specified
+        properties_to_add = new_rel.get("property", None)
 
         for rel in finalised_graph_structure["relationships"]:
             if rel["cypher_representation"] == relationship_type:
@@ -150,8 +153,16 @@ async def enrich_nodes_properties(llm, model, finalised_graph_structure):
     for new_node in enriched_node_properties["nodeProperties"]:
         label = new_node["label"]
         properties_to_add = new_node["properties"]
-
-        for node in finalised_graph_structure["nodes"]:
+        # FS-67 - Determine if the structure has nested 'nodes' or not
+        # This is a temporary code as we will change it in future to have LLM return the same structure for both Mistral
+        # and OpenAI.
+        if isinstance(finalised_graph_structure["nodes"], dict):
+            # Nodes are nested under a "nodes" key in a dictionary in the case of OpenAI
+            node_list = finalised_graph_structure["nodes"]["nodes"]
+        else:
+            # Nodes are directly in the list in the case of Mistral
+            node_list = finalised_graph_structure["nodes"]
+        for node in node_list:
             if node["label"] == label:
                 if "properties" not in node:
                     node["properties"] = []