Update app.py

app.py

@@ -55,10 +55,8 @@ Avoid assumptions. Focus on summarizing what the data reveals.
     user_prompt = f"""
 SQL Query:
 {sql_query}
-
 Query Result:
 {query_result}
-
 Explanation:
 """
     return system_prompt.strip(), user_prompt.strip()
@@ -75,7 +73,41 @@ def explain_sql_output(sql_query, query_result):
     return chat_completion.choices[0].message.content.strip()
 
 
-### --- Gradio Interface --- ###
+### --- TAB 3: SQL Follow-Ups --- ###
+def followup_prompt(sql_query, query_result):
+    system_prompt = """
+You are an assistant that suggests insightful follow-up questions based on SQL query output.
+
+Instructions:
+- Use the SQL query and result to infer possible next questions.
+- Output 3 to 5 follow-up questions.
+- Be helpful, curious, and relevant to the data.
+- Do NOT repeat the explanation. Only return follow-up questions.
+- Format as a list.
+"""
+    user_prompt = f"""
+SQL Query:
+{sql_query}
+Query Result:
+{query_result}
+
+Follow-Up Questions:
+"""
+    return system_prompt.strip(), user_prompt.strip()
+
+def generate_followups(sql_query, query_result):
+    system_prompt, user_prompt = followup_prompt(sql_query, query_result)
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt}
+        ],
+        model="llama3-70b-8192"
+    )
+    return chat_completion.choices[0].message.content.strip()
+
+
+### --- Gradio Interfaces --- ###
 tab1 = gr.Interface(
     fn=response,
     inputs=gr.JSON(label="Input JSON (question, schema)"),
@@ -95,7 +127,18 @@ tab2 = gr.Interface(
     description="Input a SQL query and its result. Get an AI-generated explanation."
 )
 
-demo = gr.TabbedInterface([tab1, tab2], ["SQL Generator", "Explain Output"])
+tab3 = gr.Interface(
+    fn=generate_followups,
+    inputs=[
+        gr.Textbox(label="SQL Query"),
+        gr.Textbox(label="SQL Output (Raw JSON or Table Result)")
+    ],
+    outputs="text",
+    title="Follow-Up Generator (Groq + LLaMA3)",
+    description="Suggests follow-up questions based on a SQL query and its result."
+)
+
+demo = gr.TabbedInterface([tab1, tab2, tab3], ["SQL Generator", "Explain Output", "SQL Follow-Ups"])
 
 if __name__ == '__main__':
     demo.launch()