Update app.py

app.py

@@ -137,6 +137,7 @@ def main():
     st.markdown('<div class="title">KASOTI</div>', unsafe_allow_html=True)
     st.markdown('<div class="subtitle">The Ultimate Guessing Game</div>', unsafe_allow_html=True)
     
+    # Initialize session state
     if 'game_state' not in st.session_state:
         st.session_state.game_state = "start"
         st.session_state.questions = []
@@ -144,8 +145,9 @@ def main():
         st.session_state.answers = []
         st.session_state.conversation_history = []
         st.session_state.category = None
-        st.session_state.waiting_for_answer = False
-    
+        st.session_state.waiting_for_answer = True
+        st.session_state.question_generated = False
+
     # Start screen
     if st.session_state.game_state == "start":
         st.markdown("""
@@ -172,21 +174,24 @@ def main():
                     st.error("Please enter either 'person', 'place', or 'object'!")
                 else:
                     st.session_state.category = category_input
-                    # Generate first question
-                    first_question = ask_llama([{
-                        "role": "user", 
-                        "content": "Ask your first yes/no question."
-                    }], category_input)
-                    st.session_state.questions = [first_question]
-                    st.session_state.conversation_history = [
-                        {"role": "assistant", "content": first_question}
-                    ]
                     st.session_state.game_state = "gameplay"
-                    st.session_state.waiting_for_answer = True
                     st.rerun()
 
-    # Gameplay screen
+    # Gameplay screen - generate first question if needed
     elif st.session_state.game_state == "gameplay":
+        if not st.session_state.questions:
+            # Generate first question
+            first_question = ask_llama([{
+                "role": "user", 
+                "content": "Ask your first yes/no question."
+            }], st.session_state.category)
+            st.session_state.questions = [first_question]
+            st.session_state.conversation_history = [
+                {"role": "assistant", "content": first_question}
+            ]
+            st.session_state.waiting_for_answer = True
+            st.rerun()
+
         # Show current question
         current_question = st.session_state.questions[st.session_state.current_q]
         st.markdown(f'<div class="question-box">Question {st.session_state.current_q + 1}/20:<br><br>'
@@ -206,24 +211,28 @@ def main():
                     st.session_state.conversation_history.append(
                         {"role": "user", "content": answer_input}
                     )
+                    st.session_state.waiting_for_answer = False
                     
                     # Check if we've reached max questions
                     if st.session_state.current_q >= 19:
                         st.session_state.game_state = "result"
-                    else:
-                        # Generate next question based on all previous Q&A
-                        next_question = ask_llama(
-                            st.session_state.conversation_history,
-                            st.session_state.category
-                        )
-                        st.session_state.questions.append(next_question)
-                        st.session_state.conversation_history.append(
-                            {"role": "assistant", "content": next_question}
-                        )
-                        st.session_state.current_q += 1
-                    
                     st.rerun()
 
+        # Generate next question if answer was submitted
+        if not st.session_state.waiting_for_answer and st.session_state.current_q < 19:
+            # Generate next question based on all previous Q&A
+            next_question = ask_llama(
+                st.session_state.conversation_history,
+                st.session_state.category
+            )
+            st.session_state.questions.append(next_question)
+            st.session_state.conversation_history.append(
+                {"role": "assistant", "content": next_question}
+            )
+            st.session_state.current_q += 1
+            st.session_state.waiting_for_answer = True
+            st.rerun()
+
         # Early guess button
         if st.button("I think I can guess now!"):
             st.session_state.game_state = "result"