Skip to content

Commit

Permalink
Changes to make OpenAI key work for dataStoreAgent (#6)
Browse files Browse the repository at this point in the history
* Changes to make OpenAI key work for dataStoreAgent

* Pushing the changes for the Rel error for mistral

* Added a comment about teh code being temporary.

* Fixing the lint error

---------
  • Loading branch information
gaganahluwalia authored Oct 23, 2024
1 parent 23a0746 commit c9b4577
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions backend/src/utils/semantic_layer_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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"] = []
Expand Down

0 comments on commit c9b4577

Please sign in to comment.