Update app.py

app.py

@@ -1494,8 +1494,10 @@ def show_story():
                 else:
                     # Regular AI Response with ending mode indicator
                     ending_info = ""
-                    if st.session_state.get('ending_mode'):
-                        ending_info = f"🎭 เหลืออีก {st.session_state.sentences_to_end} ประโยค"
+                    # แสดงจำนวนประโยคที่เหลือเฉพาะเมื่ออยู่ในโหมดจบเรื่อง
+                    if (st.session_state.get('ending_mode') and 
+                        'remaining_sentences' in entry):
+                        ending_info = f"🎭 เหลืออีก {entry['remaining_sentences']} ประโยค"
                     
                     st.markdown(f"""
                         <div style="
@@ -1521,8 +1523,10 @@ def show_story():
                 border_color = "#4caf50" if entry.get('is_correct') else "#1e88e5"
                 
                 ending_info = ""
-                if st.session_state.get('ending_mode'):
-                    ending_info = f"🎭 เหลืออีก {st.session_state.sentences_to_end} ประโยค"
+                # แสดงจำนวนประโยคที่เหลือเฉพาะเมื่ออยู่ในโหมดจบเรื่อง
+                if (st.session_state.get('ending_mode') and 
+                    'remaining_sentences' in entry):
+                    ending_info = f"🎭 เหลืออีก {entry['remaining_sentences']} ประโยค"
                 
                 st.markdown(f"""
                     <div style="
@@ -1612,20 +1616,21 @@ def show_story_ending_options():
 def handle_ending_mode(text: str):
     """Handle story submission during ending mode"""
     try:
+        # Get current remaining sentences
         remaining = st.session_state.sentences_to_end
-        
+
         # Add user's sentence
         st.session_state.story.append({
             "role": "You",
             "content": text,
             "is_corrected": False,
-            "is_correct": True,  # We don't check correctness in ending mode
-            "timestamp": datetime.now().isoformat()
+            "is_correct": True,
+            "timestamp": datetime.now().isoformat(),
+            "remaining_sentences": remaining  # เพิ่มจำนวนประโยคที่เหลือ
         })
         
-        # Decrease remaining count
+        # Decrease remaining count for user's sentence
         remaining -= 1
-        st.session_state.sentences_to_end = remaining
         
         if remaining > 0:
             # Generate AI's ending continuation
@@ -1635,14 +1640,19 @@ def handle_ending_mode(text: str):
                 remaining_sentences=remaining
             )
             
+            # Add AI response with remaining count
             st.session_state.story.append({
                 "role": "AI",
                 "content": ai_response,
-                "timestamp": datetime.now().isoformat()
+                "timestamp": datetime.now().isoformat(),
+                "remaining_sentences": remaining  # เพิ่มจำนวนประโยคที่เหลือ
             })
             
             # Decrease for AI's sentence
-            st.session_state.sentences_to_end = remaining - 1
+            remaining -= 1
+            
+            # Update remaining sentences in session state
+            st.session_state.sentences_to_end = remaining
             
         else:
             # Generate final ending
@@ -1651,6 +1661,7 @@ def handle_ending_mode(text: str):
                 st.session_state.ending_type
             )
             
+            # Add final AI response
             st.session_state.story.append({
                 "role": "AI",
                 "content": final_response,
@@ -1662,7 +1673,7 @@ def handle_ending_mode(text: str):
             st.session_state.story_completed = True
             complete_story()
         
-        # Clear input
+        # Clear input and rerun
         st.session_state.story_input_area = ""
         st.rerun()