Update chat.py

aws-samples · Aug 13, 2024 · aff91ec · aff91ec
1 parent ae6d7f2
commit aff91ec
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/backend/app/usecases/chat.py b/backend/app/usecases/chat.py
@@ -214,7 +214,10 @@ def trace_to_root(
     """Trace message map from leaf node to root node."""
     result = []
     if not node_id or node_id == "system":
-        node_id = "instruction" if "instruction" in message_map else "system"
+        if message_map:
+            node_id = list(message_map.keys())[-1]
+        else:
+            node_id = "system"
 
     current_node = message_map.get(node_id)
     while current_node: