Update app.py

app.py

@@ -233,7 +233,7 @@ def main():
     # Load custom CSS
     load_custom_css()
 
-    # Initialize session state
+    # Initialize session states
     if 'pipeline' not in st.session_state:
         st.session_state.pipeline = None
     
@@ -248,6 +248,11 @@ def main():
         st.session_state.qa_history_loaded = True
         load_qa_history()
 
+    # Initialize pipeline
+    if st.session_state.pipeline is None:
+        with st.spinner("กำลังเริ่มต้นระบบ..."):
+            st.session_state.pipeline = initialize_pipeline()
+
     # Header
     st.markdown("""
         <div style="text-align: center; padding: 2rem 0;">
@@ -256,59 +261,41 @@ def main():
         </div>
     """, unsafe_allow_html=True)
 
-    # Initialize pipeline
-    if st.session_state.pipeline is None:
-        with st.spinner("กำลังเริ่มต้นระบบ..."):
-            st.session_state.pipeline = initialize_pipeline()
+    chat_col, info_col = st.columns([7, 3])
 
     def clear_text():
-        st.session_state.query = ''
         st.session_state.query_input = ''
     
     def handle_submit():
-        # Get the current query value
         query = st.session_state.query_input
         if query:
             if st.session_state.pipeline is None:
                 st.error("❌ ไม่สามารถเชื่อมต่อกับระบบได้ กรุณาลองใหม่อีกครั้ง")
                 return
                 
-            add_to_history("user", query)
-                
             try:
                 with st.spinner("🔍 กำลังค้นหาคำตอบ..."):
                     result = st.session_state.pipeline.process_query(query)
-                    add_to_history("assistant", result["answer"])
-                    
-                    with st.expander("📚 แสดงข้อมูลอ้างอิง", expanded=False):
-                        for i, doc in enumerate(result["documents"], 1):
-                            st.markdown(f"""
-                                <div style="padding: 1rem; background-color: #F9FAFB; border-radius: 8px; margin: 0.5rem 0;">
-                                    <strong>เอกสารที่ {i}:</strong><br>
-                                    {doc.content}
-                                </div>
-                            """, unsafe_allow_html=True)
                     
-                    with st.expander("🔍 รายละเอียดการวิเคราะห์คำถาม", expanded=False):
-                        st.json(result["query_info"])
+                    # Update session state
+                    st.session_state.chat_history.append(("user", query))
+                    st.session_state.chat_history.append(("assistant", {
+                        "answer": result["answer"],
+                        "documents": result["documents"],
+                        "query_info": result["query_info"]
+                    }))
                     
-                    # Clear the input
+                    # Clear input
                     clear_text()
-                    st.rerun()
                     
             except Exception as e:
                 st.error(f"❌ เกิดข้อผิดพลาด: {str(e)}")
         else:
             st.warning("⚠️ กรุณาระบุคำถาม")
 
-    chat_col, info_col = st.columns([7, 3])
-    
     with chat_col:
+        # Input section at the top
         with st.container():
-            # Display chat history
-            display_chat_history()
-            
-            # Query input section
             st.markdown("""
                 <label for="query_input" style="font-size: 1.2rem; font-weight: 600; margin-bottom: 1rem; display: block;">
                     <span style="color: #ffffff; border-left: 4px solid #ffffff; padding-left: 0.8rem;">
@@ -317,35 +304,57 @@ def main():
                 </label>
             """, unsafe_allow_html=True)
             
-            # Create the text input with callbacks
+            # Text input
             query = st.text_input(
                 "",
                 key="query_input",
                 placeholder="เช่น: วันสุดท้ายของการสอบปากเปล่าในภาคเรียนที่ 1/2567 คือวันที่เท่าไร?",
-                on_change=handle_submit
             )
 
             # Button layout
             col1, col2, col3 = st.columns([1, 1, 4])
-        
+            
             with col1:
-                st.button(
-                    "📤 ส่งคำถาม",
-                    type="primary",
-                    use_container_width=True,
-                    on_click=handle_submit
-                )
-        
+                if st.button("📤 ส่งคำถาม", type="primary", use_container_width=True):
+                    handle_submit()
+            
             with col2:
-                st.button(
-                    "🗑️ ล้างประวัติ",
-                    type="secondary",
-                    use_container_width=True,
-                    on_click=lambda: setattr(st.session_state, 'chat_history', []) or st.rerun()
-                )
+                if st.button("🗑️ ล้างประวัติ", type="secondary", use_container_width=True):
+                    st.session_state.chat_history = []
+                    st.rerun()
+
+        # Chat history display
+        for i, (role, content) in enumerate(st.session_state.chat_history):
+            if role == "user":
+                st.markdown(f"""
+                    <div class="chat-message user-message">
+                        <strong>🧑 คำถาม:</strong><br>
+                        {content}
+                    </div>
+                """, unsafe_allow_html=True)
+            else:
+                st.markdown(f"""
+                    <div class="chat-message assistant-message">
+                        <strong>🤖 คำตอบ:</strong><br>
+                        {content['answer']}
+                    </div>
+                """, unsafe_allow_html=True)
+                
+                # Add expanders right after the answer
+                with st.expander("📚 แสดงข้อมูลอ้างอิง", expanded=False):
+                    for i, doc in enumerate(content['documents'], 1):
+                        st.markdown(f"""
+                            <div style="padding: 1rem; background-color: #F9FAFB; border-radius: 8px; margin: 0.5rem 0;">
+                                <strong>เอกสารที่ {i}:</strong><br>
+                                {doc.content}
+                            </div>
+                        """, unsafe_allow_html=True)
+                
+                with st.expander("🔍 รายละเอียดการวิเคราะห์คำถาม", expanded=False):
+                    st.json(content['query_info'])
 
+    # Info column content remains the same
     with info_col:
-        # System information
         st.markdown("""
             <div style="background-color: #F9FAFB; padding: 1.5rem; border-radius: 12px; margin-bottom: 2rem;">
                 <h3 style="color: #1E3A8A;">ℹ️ เกี่ยวกับระบบ</h3>
@@ -364,7 +373,6 @@ def main():
             </div>
         """, unsafe_allow_html=True)
         
-        # System status
         st.markdown("""
             <div style="background-color: #f9fafb; padding: 1.5rem; border-radius: 12px;">
                 <h3 style="color: #1E3A8A;">🔄 สถานะระบบ</h3>