Update app.py

app.py

@@ -364,11 +364,16 @@ def main():
                     </div>
                 """, unsafe_allow_html=True)
             else:
-                if isinstance(content, dict):  
+                # Handle both string and dictionary responses
+                assistant_response = content
+                reference_docs = None
+                
+                # Check if the response is a dictionary with both answer and documents
+                if isinstance(content, dict):
                     assistant_response = content.get('answer', '❌ ไม่มีข้อมูลคำตอบ')
-                else:  
-                    assistant_response = content
+                    reference_docs = content.get('documents', [])
 
+                # Display the assistant's response
                 st.markdown(f"""
                     <div class="chat-message assistant-message">
                         <strong>🤖 คำตอบ:</strong><br>
@@ -376,16 +381,18 @@ def main():
                     </div>
                 """, unsafe_allow_html=True)
 
-                if isinstance(content, dict) and content.get('documents'):
+                # Display reference documents if available
+                if reference_docs:
                     with st.expander("📚 แสดงข้อมูลอ้างอิง", expanded=False):
-                        for i, doc in enumerate(content['documents'], 1):
+                        for idx, doc in enumerate(reference_docs, 1):
+                            doc_content = doc.get('content', '') if isinstance(doc, dict) else str(doc)
                             st.markdown(f"""
-                                <div style="padding: 1rem; background-color: #000000; border-radius: 8px; margin: 0.5rem 0;">
-                                    <strong>เอกสารที่ {i}:</strong><br>
-                                    {doc.content}
+                                <div style="padding: 1rem; background-color: #F3F4F6; border-radius: 8px; margin: 0.5rem 0; border: 1px solid #E5E7EB;">
+                                    <strong>เอกสารที่ {idx}:</strong><br>
+                                    {doc_content}
                                 </div>
                             """, unsafe_allow_html=True)
-        
+
         # Create chat input using the new implementation
         create_chat_input()