Update utils.py

utils.py

@@ -1,4 +1,5 @@
 import pandas as pd
+from ast import literal_eval
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
@@ -116,7 +117,24 @@ def extract_table_name(llm, query):
     for table in metadata_df.table.unique():
         if table in pred:
             return table
-            
+
+def extract_question_list(llm, query):
+    sys_prompt = """You are an AI assistant that determines if multiple questions are stacked in a single question and split the question into sub questions and return a list of them. Make sure the response is a valid Python list. 
+    If the question is a single question and return the original question."""
+
+    message1 = SystemMessage(content=sys_prompt)
+
+    message2 = HumanMessage(
+        content=query
+    )
+    message_log = [message1, message2]
+    
+    output = llm.invoke(message_log)
+    pred = output.content
+    
+    return pred
+
+
 warnings.filterwarnings('ignore', message="pandas only supports SQLAlchemy connectable.*", category=UserWarning, module='chain')
 
 intermediate_steps_KEY = "intermediate_steps"