Update app.py

app.py

@@ -119,13 +119,15 @@ def ask_llama(conversation_history, category, is_final_guess=False):
     }
 
     system_prompt = f"""You're playing 20 questions to guess a {category}. Follow these rules:
-1. Ask strategic, non-repeating yes/no questions that narrow down possibilities
-2. Consider all previous answers carefully before asking next question
-3. If you're very confident (80%+ sure), respond with "Final Guess: [your guess]"
-4. For places: ask about continent, climate, famous landmarks, country, city or population
-5. For people: ask about fictional or real, profession, gender, alive/dead, nationality, or fame
-6. For objects: ask about size, color, usage, material, or where it's found
-7. Never repeat questions and always make progress toward guessing"""
+1. Ask strategic, non-repeating yes/no questions that narrow down possibilities.
+2. Consider all previous answers carefully before asking the next question.
+3. If you're very confident (80%+ sure), respond with "Final Guess: [your guess]".
+4. For places: ask about continent, climate, famous landmarks, country, city or population.
+5. For people: ask about fictional or real, profession, gender, alive/dead, nationality, or fame.
+6. For objects: ask about size, color, usage, material, or where it's found.
+7. Never repeat questions and always make progress toward guessing.
+8. If the guess is generic (e.g. "bed") and the player confirms it with "yes", end the game.
+9. If the guess is generic but the answer is "no", ask additional questions to refine the guess."""
 
     if is_final_guess:
         prompt = f"""Based on these answers about a {category}, provide ONLY your final guess with no extra text:
@@ -214,13 +216,14 @@ def main():
                         {"role": "assistant", "content": first_question}
                     ]
                     st.session_state.game_state = "gameplay"
+                    st.session_state.current_q = 0  # start at question 0
                     st.rerun()
 
     # Gameplay screen
     elif st.session_state.game_state == "gameplay":
         current_question = st.session_state.questions[st.session_state.current_q]
         
-        # Check if AI made a guess
+        # If the AI made a final guess, switch to confirmation
         if "Final Guess:" in current_question:
             st.session_state.final_guess = current_question.split("Final Guess:")[1].strip()
             st.session_state.game_state = "confirm_guess"
@@ -242,13 +245,12 @@ def main():
                         {"role": "user", "content": answer_input}
                     )
 
-                    # Generate next response
+                    # Generate next response from Llama
                     next_response = ask_llama(
                         st.session_state.conversation_history,
                         st.session_state.category
                     )
                     
-                    # Check if AI made a guess
                     if "Final Guess:" in next_response:
                         st.session_state.final_guess = next_response.split("Final Guess:")[1].strip()
                         st.session_state.game_state = "confirm_guess"
@@ -279,7 +281,7 @@ def main():
                     st.markdown(f"**You:** {msg['query']}")
                     st.markdown(f"**Help Assistant:** {msg['response']}")
 
-    # Guess confirmation screen using text input response
+    # Confirm guess screen
     elif st.session_state.game_state == "confirm_guess":
         st.markdown(f'<div class="question-box">🤖 My Final Guess:<br><br>'
                     f'<strong>Is it {st.session_state.final_guess}?</strong></div>',
@@ -294,9 +296,8 @@ def main():
                     if confirm_input == "yes":
                         st.session_state.game_state = "result"
                         st.rerun()
-                        st.stop()  # Immediately halt further execution
                     else:
-                        # Add negative response to history and continue gameplay
+                        # Add the negative response to history and continue gameplay
                         st.session_state.conversation_history.append(
                             {"role": "user", "content": "no"}
                         )
@@ -333,7 +334,7 @@ def main():
         time.sleep(1)
         st.markdown(f'<div class="final-reveal" style="font-size:3.5rem;color:#6C63FF;">{st.session_state.final_guess}</div>',
                     unsafe_allow_html=True)
-        st.markdown(f"<p style='text-align:center'>Guessed in {len(st.session_state.questions)} questions</p>", 
+        st.markdown(f"<p style='text-align:center'>Guessed in {st.session_state.current_q + 1} questions</p>", 
                     unsafe_allow_html=True)
         
         if st.button("Play Again", key="play_again"):