Update app.py

app.py

@@ -549,6 +549,14 @@ def init_session_state():
         st.session_state.clear_input = False
     if 'text_input' not in st.session_state:
         st.session_state.text_input = ""
+    if 'ending_mode' not in st.session_state:
+        st.session_state.ending_mode = False
+    if 'sentences_to_end' not in st.session_state:
+        st.session_state.sentences_to_end = 0
+    if 'ending_type' not in st.session_state:
+        st.session_state.ending_type = None
+    if 'story_completed' not in st.session_state:
+        st.session_state.story_completed = False
 
 def init_theme_state():
     """Initialize theme-specific state variables"""
@@ -1626,11 +1634,12 @@ def handle_ending_mode(text: str):
             "is_corrected": False,
             "is_correct": True,
             "timestamp": datetime.now().isoformat(),
-            "remaining_sentences": remaining  # เพิ่มจำนวนประโยคที่เหลือ
+            "remaining_sentences": remaining  # Include remaining sentences
         })
         
         # Decrease remaining count for user's sentence
-        remaining -= 1
+        st.session_state.sentences_to_end -= 1
+        remaining = st.session_state.sentences_to_end  # Update local variable
         
         if remaining > 0:
             # Generate AI's ending continuation
@@ -1645,15 +1654,12 @@ def handle_ending_mode(text: str):
                 "role": "AI",
                 "content": ai_response,
                 "timestamp": datetime.now().isoformat(),
-                "remaining_sentences": remaining  # เพิ่มจำนวนประโยคที่เหลือ
+                "remaining_sentences": remaining
             })
             
-            # Decrease for AI's sentence
-            remaining -= 1
-            
-            # Update remaining sentences in session state
-            st.session_state.sentences_to_end = remaining
-            
+            # Decrease remaining sentences
+            st.session_state.sentences_to_end -= 1
+            remaining = st.session_state.sentences_to_end
         else:
             # Generate final ending
             final_response = generate_final_ending(
@@ -1674,13 +1680,15 @@ def handle_ending_mode(text: str):
             complete_story()
         
         # Clear input and rerun
-        st.session_state.story_input_area = ""
+        st.session_state.text_input = ""
+        st.session_state.clear_input = True
         st.rerun()
         
     except Exception as e:
         logging.error(f"Error in ending mode: {str(e)}")
         st.error("เกิดข้อผิดพลาดในโหมดจบเรื่อง กรุณาลองใหม่อีกครั้ง")
 
+
 def generate_final_ending(story: List[dict], ending_type: str) -> str:
     """Generate the final ending sentence based on the story and chosen ending type"""
     try:
@@ -1756,7 +1764,7 @@ def generate_ending_continuation(text: str, ending_type: str, remaining_sentence
                 - Build towards joy, success, or resolution
                 - Use uplifting and positive language
                 - Connect to previous story elements
-                """,
+            """,
             "Mysterious Ending": """
                 Role: Mystery writer creating intrigue
                 Goal: Leave readers thinking and wondering
@@ -1764,7 +1772,7 @@ def generate_ending_continuation(text: str, ending_type: str, remaining_sentence
                 - Add subtle hints and clues
                 - Create atmospheric descriptions
                 - Leave some questions unanswered
-                """,
+            """,
             "Lesson Learned": """
                 Role: Moral story concluder
                 Goal: Incorporate meaningful life lessons
@@ -1772,7 +1780,7 @@ def generate_ending_continuation(text: str, ending_type: str, remaining_sentence
                 - Connect actions to consequences
                 - Show character growth
                 - Express the moral naturally
-                """,
+            """,
             "Surprise Ending": """
                 Role: Plot twist creator
                 Goal: Deliver unexpected but satisfying conclusion
@@ -1780,11 +1788,11 @@ def generate_ending_continuation(text: str, ending_type: str, remaining_sentence
                 - Plant subtle hints earlier
                 - Subvert expectations logically
                 - Maintain story coherence
-                """
+            """
         }
 
         response = client.chat.completions.create(
-            model="gpt-4",
+            model="gpt-4o-mini",
             messages=[
                 {
                     "role": "system",
@@ -1792,10 +1800,11 @@ def generate_ending_continuation(text: str, ending_type: str, remaining_sentence
                     {ending_prompts[ending_type]}
                     
                     Additional Rules:
-                    - You have {remaining_sentences} sentences to conclude
-                    - Each response should be max 2 sentences
-                    - Build towards the finale naturally
-                    - Connect to previous story elements
+                    - You have {remaining_sentences} sentences left to conclude the story.
+                    - Each response should be maximum 2 sentences.
+                    - Build towards the finale naturally.
+                    - Ensure coherence with previous story events.
+                    - Indicate awareness of the remaining sentences.
                     """
                 },
                 {
@@ -1806,13 +1815,14 @@ def generate_ending_continuation(text: str, ending_type: str, remaining_sentence
             max_tokens=100,
             temperature=0.7
         )
-        
+
         return response.choices[0].message.content.strip()
         
     except Exception as e:
-        logging.error(f"Error generating ending: {str(e)}")
+        logging.error(f"Error generating ending continuation: {str(e)}")
         return "The story moved towards its conclusion..."
 
+
 def complete_story():
     """Handle story completion and celebration"""
     st.balloons()  # แสดงเอฟเฟคฉลอง
@@ -2173,11 +2183,16 @@ def show_story_input():
     # If clear_input flag is True, reset it
     if st.session_state.clear_input:
         st.session_state.clear_input = False
+
+    # Show remaining sentences if in ending mode
+    if st.session_state.get('ending_mode'):
+        remaining = st.session_state.sentences_to_end
+        st.info(f"🎭 โหมดจบเรื่อง - เหลืออีก {remaining} ประโยค")
     
     # Input area
     text_input = st.text_area(
         "เขียนต่อจากเรื่องราว | Continue the story:",
-        value=default_value,
+        value=st.session_state.get('text_input', ""),
         height=100,
         key="story_input_area",
         help="พิมพ์ประโยคภาษาอังกฤษเพื่อต่อเรื่อง",
@@ -2194,7 +2209,6 @@ def show_story_input():
             if not text_input.strip():
                 st.warning("กรุณาเขียนข้อความก่อนส่ง")
                 return
-                
             try:
                 with st.spinner("กำลังวิเคราะห์ประโยค..."):
                     handle_story_submission(text_input.strip())
@@ -2216,6 +2230,12 @@ def handle_story_submission(text: str):
         st.error("กรุณาเลือกธีมเรื่องราวก่อนเริ่มเขียน")
         return
 
+    # Check if in ending mode
+    if st.session_state.get('ending_mode'):
+        handle_ending_mode(text)
+        return
+
+    # Regular processing
     try:
         # Get feedback
         feedback_data = provide_feedback(text, st.session_state.level)
@@ -2242,7 +2262,7 @@ def handle_story_submission(text: str):
         # Generate AI continuation
         try:
             logging.info("Attempting to generate AI continuation...")
-            # ใช้ข้อความที่ถูกต้องสำหรับการต่อเรื่อง
+            # Use the corrected text if there are errors
             text_for_continuation = feedback_data['corrected'] if feedback_data.get('has_errors') else text
             
             ai_response = generate_story_continuation(text_for_continuation, st.session_state.level)
@@ -2259,7 +2279,7 @@ def handle_story_submission(text: str):
                 logging.info("AI response added to story successfully")
             else:
                 logging.error("AI generated empty response")
-                # กรณีที่ AI ไม่สร้างประโยค ให้สร้างประโยคง่ายๆ ตาม theme
+                # Create a fallback response
                 fallback_response = generate_fallback_response(st.session_state.current_theme, st.session_state.level)
                 st.session_state.story.append({
                     "role": "AI",
@@ -2268,7 +2288,7 @@ def handle_story_submission(text: str):
                 })
         except Exception as e:
             logging.error(f"Error generating AI continuation: {str(e)}")
-            # สร้าง fallback response เมื่อเกิดข้อผิดพลาด
+            # Create a fallback response
             fallback_response = generate_fallback_response(st.session_state.current_theme, st.session_state.level)
             st.session_state.story.append({
                 "role": "AI",
@@ -2289,6 +2309,7 @@ def handle_story_submission(text: str):
         logging.error(f"Error in story submission: {str(e)}")
         raise
 
+
 def generate_fallback_response(theme_id: str, level: str) -> str:
     """Generate a simple fallback response when AI continuation fails"""
     try:
@@ -2362,9 +2383,11 @@ def main():
         if 'ending_mode' not in st.session_state:
             st.session_state.ending_mode = False
         if 'sentences_to_end' not in st.session_state:
-            st.session_state.sentences_to_end = 5
+            st.session_state.sentences_to_end = 0  # Set to 0 by default
         if 'ending_type' not in st.session_state:
             st.session_state.ending_type = None
+        if 'story_completed' not in st.session_state:
+            st.session_state.story_completed = False
         
         # Add watermark
         st.markdown("""
@@ -2394,7 +2417,7 @@ def main():
             else:
                 # Show story progress if story exists
                 if st.session_state.story:
-                    # แสดงความคืบหน้าของเรื่อง
+                    # Display story progress
                     total_sentences = len(st.session_state.story)
                     st.markdown(f"""
                         <div style="
@@ -2407,7 +2430,7 @@ def main():
                                 📊 ความยาวเรื่อง: {total_sentences} ประโยค
                             </div>
                             <div class="progress-bar">
-                                <div class="progress-bar-fill" style="width: {min(total_sentences * 10, 100)}%;"></div>
+                                <div class="progress-bar-fill" style="width: {min(total_sentences * 5, 100)}%;"></div>
                             </div>
                             <div style="font-size: 0.9em; color: #666; margin-top: 5px;">
                                 เรื่องควรยาว 10-20 ประโยค เพื่อความสมบูรณ์
@@ -2415,16 +2438,16 @@ def main():
                         </div>
                     """, unsafe_allow_html=True)
                     
-                    # แสดงตัวเลือกจบเรื่องเมื่อเรื่องยาวพอ
-                    if (total_sentences >= 10 and not st.session_state.get('ending_mode')):
+                    # Show ending options if story is long enough and ending mode is not active
+                    if (total_sentences >= 10 and not st.session_state.get('ending_mode') and not st.session_state.get('story_completed')):
                         st.markdown("### 🎭 ต้องการจบเรื่องหรือไม่?")
                         ending_type = st.radio(
                             "เลือกวิธีจบเรื่อง:",
                             options=[
-                                "Happy Ending - จบแบบมีความสุข",
-                                "Mysterious Ending - จบแบบทิ้งท้ายให้คิดต่อ",
-                                "Lesson Learned - จบแบบได้ข้อคิด",
-                                "Surprise Ending - จบแบบพลิกความคาดหมาย"
+                                "Happy Ending",
+                                "Mysterious Ending",
+                                "Lesson Learned",
+                                "Surprise Ending"
                             ],
                             index=0,
                             key="ending_type_selector"
@@ -2433,10 +2456,11 @@ def main():
                         if st.button("🎬 เริ่มจบเรื่อง", use_container_width=True):
                             st.session_state.ending_mode = True
                             st.session_state.ending_type = ending_type
-                            st.session_state.sentences_to_end = 5
-                            st.rerun()
+                            st.session_state.sentences_to_end = 5  # Set the countdown
+                            st.success(f"โหมดจบเรื่องเริ่มต้นแล้ว! รูปแบบการจบ: {ending_type}")
+                            st.experimental_rerun()
                     
-                    # แสดงแจ้งเตือนโหมดจบเรื่อง
+                    # Show ending mode notification
                     if st.session_state.get('ending_mode'):
                         remaining = st.session_state.sentences_to_end
                         if remaining > 0:
@@ -2447,36 +2471,50 @@ def main():
                                 
                                 พยายามเขียนให้เรื่องจบอย่างสมบูรณ์!
                             """)
-                
-                # แสดงส่วนหลัก
-                show_main_interface()
-                
-                # ตรวจสอบการจบเรื่อง
-                if st.session_state.get('ending_mode'):
-                    if st.session_state.sentences_to_end <= 0:
-                        if not st.session_state.get('story_completed'):
-                            complete_story()
+                        elif remaining <= 0 and not st.session_state.get('story_completed'):
+                            # Story should be marked as completed
+                            st.success("🎉 เรื่องราวของคุณจบลงแล้ว!")
                             st.balloons()
-                    elif st.session_state.sentences_to_end == 1:
-                        st.info("⚡ เหลือประโยคสุดท้ายแล้ว! เขียนให้จบอย่างสวยงาม")
+                            st.session_state.ending_mode = False
+                            st.session_state.story_completed = True
 
+                # Display main interface
+                show_main_interface()
+                
+                # Handle story completion
+                if st.session_state.get('story_completed'):
+                    st.markdown("""
+                        <div style="
+                            background-color: #e8f5e9;
+                            padding: 20px;
+                            border-radius: 10px;
+                            text-align: center;
+                            margin: 20px 0;
+                        ">
+                            <h2>🎉 ยินดีด้วย! คุณเขียนเรื่องราวจบสมบูรณ์แล้ว</h2>
+                            <p>คุณสามารถบันทึกเรื่องราวหรือเริ่มเรื่องใหม่ได้</p>
+                        </div>
+                    """, unsafe_allow_html=True)
+                    # Show save and reset options
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        if st.button("💾 บันทึกเรื่องราว", use_container_width=True):
+                            save_completed_story()
+                    with col2:
+                        if st.button("🔄 เริ่มเรื่องใหม่", use_container_width=True):
+                            reset_story()
+                            st.experimental_rerun()
+                
         # Handle reset if needed
         if st.session_state.get('should_reset'):
             reset_story()
             # Reset ending states
             st.session_state.ending_mode = False
-            st.session_state.sentences_to_end = 5
+            st.session_state.sentences_to_end = 0
             st.session_state.ending_type = None
             st.session_state.story_completed = False
-
-        # Handle reset if needed
-        if st.session_state.should_reset:
-            reset_story()
-            # Reset ending mode states
-            st.session_state.ending_mode = False
-            st.session_state.sentences_to_end = 5
-            st.session_state.ending_type = None
-            
+            st.experimental_rerun()
+        
         # Auto-save progress periodically
         if st.session_state.story:
             auto_save_progress()
@@ -2484,6 +2522,7 @@ def main():
     except Exception as e:
         handle_application_error(e)
 
+
 def check_session_status():
     """Check and maintain session status"""
     try: