Update app.py

app.py

@@ -29,6 +29,8 @@ def init_session_state():
         st.session_state.total_words = 0
     if 'badges' not in st.session_state:
         st.session_state.badges = []
+    if 'should_reset' not in st.session_state:
+        st.session_state.should_reset = False
 
 init_session_state()
 
@@ -120,6 +122,19 @@ def update_achievements(text: str):
             st.session_state.badges.append(badge)
             st.success(f"🏆 Achievement Unlocked: {badge}!")
 
+def reset_story():
+    """Reset the story and related state variables."""
+    st.session_state.story = []
+    st.session_state.achievements = []
+    st.session_state.unique_words = set()
+    st.session_state.total_words = 0
+    st.session_state.badges = []
+    st.session_state.should_reset = False
+
+# Handle story reset if needed
+if st.session_state.should_reset:
+    reset_story()
+
 # Main UI Layout
 st.title("📖 JoyStory - Interactive Story Adventure")
 
@@ -133,8 +148,8 @@ with st.sidebar:
     st.session_state.level = level
     
     if st.button("Start New Story"):
-        st.session_state.story = []
-        st.experimental_rerun()
+        st.session_state.should_reset = True
+        st.rerun()
 
 # Main content area
 col1, col2 = st.columns([3, 1])
@@ -171,7 +186,8 @@ with col1:
             ai_response = generate_story_continuation(user_input, st.session_state.level)
             st.session_state.story.append({"role": "AI", "content": ai_response})
             
-            st.experimental_rerun()
+            # Clear the input and rerun
+            st.rerun()
         else:
             st.warning("Please write something to continue the story.")