Update src/streamlit_app.py

src/streamlit_app.py

@@ -56,6 +56,11 @@ p_application = """
     6. Do not include any explanation or examples in your response.
 """
 
+# count the number of times "I don't know is clicked"
+if "retry_count" not in st.session_state:
+    st.session_state.retry_count = 0
+
+    
 # Initialize session state
 if "help_clicks" not in st.session_state:
     st.session_state.help_clicks = 0
@@ -102,48 +107,63 @@ with st.container():
 
             with col1:
                 if st.button("📝 Explain the question"):
-                    prompt = (
-                        "[INST]<<SYS>>\n"
-                        f"{prompt + p_explanation}\n"
-                        "<</SYS>>\n\n"
-                        f"{QUESTION}\n"
-                        "[/INST]"
-                    )
-                    st.session_state.response = llm.invoke(prompt)
-                    
-                    st.session_state.button_clicked = "Explain the question"
+                    if st.session_state.button_clicked != "Explain the question":
+                        # First time clicked
+                        full_prompt = (
+                            "[INST]<<SYS>>\n"
+                            f"{prompt + p_explanation}\n"
+                            "<</SYS>>\n\n"
+                            f"{QUESTION}\n"
+                            "[/INST]"
+                        )
+                        st.session_state.response = llm.invoke(full_prompt)
+                        st.session_state.retry_count = 0
+                        st.session_state.full_prompt = full_prompt  # Save prompt for retry
+                        st.session_state.button_clicked = "Explain the question"
             with col2:
-                if st.button("💡 Give an example"):
-                    prompt = (
-                        "[INST]<<SYS>>\n"
-                        f"{prompt + p_example}\n"
-                        "<</SYS>>\n\n"
-                        f"{QUESTION}\n"
-                        "[/INST]"
-                    )
-                    st.session_state.response = llm.invoke(prompt)
-                    
-                    st.session_state.button_clicked = "Give an example"
+                    if st.button("💡 Give an example"):
+                        if st.session_state.button_clicked != "Give an example":
+                            # First time clicked
+                            full_prompt = (
+                                "[INST]<<SYS>>\n"
+                                f"{prompt + p_example}\n"
+                                "<</SYS>>\n\n"
+                                f"{QUESTION}\n"
+                                "[/INST]"
+                            )
+                            st.session_state.response = llm.invoke(full_prompt)
+                            st.session_state.retry_count = 0
+                            st.session_state.full_prompt = full_prompt  # Save prompt for retry
+                            st.session_state.button_clicked = "Give an example"
             with col3:
                 if st.button("🤔 Who cares?"):
-                    prompt = (
-                        "[INST]<<SYS>>\n"
-                        f"{prompt + p_application}\n"
-                        "<</SYS>>\n\n"
-                        f"{QUESTION}\n"
-                        "[/INST]"
-                    )
-                    st.session_state.response = llm.invoke(prompt)
-                    
-                    st.session_state.button_clicked = "Who cares?"
-
+                        if st.session_state.button_clicked != "Who cares?":
+                            # First time clicked
+                            full_prompt = (
+                                "[INST]<<SYS>>\n"
+                                f"{prompt + p_application}\n"
+                                "<</SYS>>\n\n"
+                                f"{QUESTION}\n"
+                                "[/INST]"
+                            )
+                            st.session_state.response = llm.invoke(full_prompt)
+                            st.session_state.retry_count = 0
+                            st.session_state.full_prompt = full_prompt  # Save prompt for retry
+                            st.session_state.button_clicked = "Who cares?"
             st.markdown("---")
 
         # Display response text if a sub-button is clicked
         if st.session_state.button_clicked:
             with st.container():
                 st.info(st.session_state.response)
-
+        
+                if st.session_state.button_clicked == "Explain the question":
+                    if st.button("I don't understand. Try again."):
+                        st.session_state.retry_count += 1
+                        alt_llm = get_llm(temperature=0.9)
+                        st.session_state.response = alt_llm.invoke(
+                            st.session_state.full_prompt
+                        )
 # Optional: Add footer or spacing
 st.markdown("<br><br>", unsafe_allow_html=True)