Update app.py

app.py

@@ -78,7 +78,8 @@ suggested_drugs = [
 # Function to handle drug query
 def query_drug(drug_name, chat_history):
     if not drug_name.strip():
-        return chat_history + [[None, "Please enter a drug name."]]
+        chat_history.append({"role": "assistant", "content": "Please enter a drug name."})
+        return chat_history
     
     try:
         result = qa_chain.invoke({"query": drug_name})
@@ -96,7 +97,13 @@ def query_drug(drug_name, chat_history):
         else:
             response = f"Error: An unexpected error occurred: {error_msg}"
     
-    return chat_history + [[drug_name, response]]
+    chat_history.append({"role": "user", "content": drug_name})
+    chat_history.append({"role": "assistant", "content": response})
+    return chat_history
+
+# Function to handle suggested drug buttons
+def query_suggested_drug(drug_name, chat_history):
+    return query_drug(drug_name, chat_history)
 
 # Gradio Interface
 with gr.Blocks(title="DrugScan") as demo:
@@ -108,26 +115,28 @@ with gr.Blocks(title="DrugScan") as demo:
     gr.Image(logo_url, width=150)
     
     # Chat interface
-    chatbot = gr.Chatbot(label="Results")
+    chatbot = gr.Chatbot(label="Results", type="messages")
     drug_input = gr.Textbox(placeholder="Enter a drug name (e.g., 'Azirox')", label="Drug Name")
     
+    # Search button
+    search_btn = gr.Button("Search")
+    search_btn.click(
+        fn=query_drug,
+        inputs=[drug_input, chatbot],
+        outputs=chatbot
+    )
+    
     # Suggested drugs buttons
     gr.Markdown("### Try These Drugs")
     with gr.Row():
         for drug in suggested_drugs:
-            gr.Button(drug).click(
-                fn=query_drug,
-                inputs=[drug, chatbot],
+            btn = gr.Button(drug)
+            btn.click(
+                fn=query_suggested_drug,
+                inputs=[gr.State(value=drug), chatbot],
                 outputs=chatbot
             )
     
-    # Search button
-    drug_input.submit(
-        fn=query_drug,
-        inputs=[drug_input, chatbot],
-        outputs=chatbot
-    )
-    
     # Disclaimer
     gr.Markdown("### Important Disclaimer")
     gr.Markdown(