Update components/chat.py

components/chat.py

@@ -4,29 +4,71 @@ from utils.database import verify_vector_store
 
 def clean_response_content(response_content):
     """Clean the response content of any technical metadata."""
-    # Convert to string if it's an AIMessage
-    response_content = str(response_content)
+    # Convert the response to string first
+    content = str(response_content)
     
-    # Remove the content= prefix if it exists
-    if "content='" in response_content:
-        response_content = response_content.split("content='")[1]
-    
-    # Remove response metadata
-    if "response_metadata=" in response_content:
-        response_content = response_content.split("response_metadata=")[0]
-    
-    # Remove any trailing quotes or metadata artifacts
-    response_content = response_content.rstrip("' ")
-    
-    # Replace escaped newlines with actual newlines
-    response_content = response_content.replace('\\n', '\n')
-    
-    return response_content
+    # Clean up the content
+    if "content='" in content:
+        content = content.split("content='")[1].split("'")[0]
+    elif 'content="' in content:
+        content = content.split('content="')[1].split('"')[0]
+        
+    return content.replace('\\n', '\n')
 
 def display_chat_interface():
     """Display modern chat interface with clean formatting."""
     
-    # [Previous CSS styling remains the same...]
+    # Add custom CSS for modern chat styling
+    st.markdown("""
+    <style>
+        /* Clean, modern chat container */
+        .chat-container {
+            max-width: 800px;
+            margin: auto;
+        }
+        
+        /* Message styling */
+        .user-message {
+            background-color: #f0f2f6;
+            padding: 1rem;
+            border-radius: 10px;
+            margin: 1rem 0;
+        }
+        
+        .assistant-message {
+            background-color: #ffffff;
+            border: 1px solid #e0e0e0;
+            padding: 1.5rem;
+            border-radius: 10px;
+            margin: 1rem 0;
+        }
+        
+        /* Section styling */
+        .section-header {
+            font-weight: bold;
+            color: #0f52ba;
+            margin-top: 1rem;
+        }
+        
+        .source-info {
+            background-color: #f8f9fa;
+            padding: 0.5rem;
+            border-left: 3px solid #0f52ba;
+            margin-top: 1rem;
+            font-size: 0.9rem;
+        }
+        
+        /* Content sections */
+        .content-section {
+            margin: 0.5rem 0;
+        }
+        
+        /* Clean bullets */
+        .clean-bullet {
+            margin: 0.3rem 0;
+        }
+    </style>
+    """, unsafe_allow_html=True)
     
     # Check if QA system is initialized
     if 'qa_system' not in st.session_state or st.session_state.qa_system is None:
@@ -46,32 +88,13 @@ def display_chat_interface():
             </div>
             """, unsafe_allow_html=True)
         elif isinstance(message, AIMessage):
-            # Clean and structure the response
-            response_content = clean_response_content(message.content)
-            
-            # Format the content into sections
-            formatted_content = ""
-            sections = response_content.split('\n\n')
-            
-            for section in sections:
-                if section.startswith('Summary:'):
-                    formatted_content += f'<div class="summary-section">{section}</div>'
-                elif section.startswith('Key Points:'):
-                    formatted_content += '<div class="key-points">'
-                    formatted_content += f'<div class="section-header">{section.split(":")[0]}:</div>'
-                    points = [p.strip() for p in section.split('\n')[1:] if p.strip()]
-                    for point in points:
-                        formatted_content += f'<div class="key-point-item">{point.lstrip("- ")}</div>'
-                    formatted_content += '</div>'
-                elif section.startswith('Source'):
-                    formatted_content += f'<div class="source-info">{section}</div>'
-                else:
-                    formatted_content += f'<div class="content-section">{section}</div>'
+            # Get the content and clean it
+            content = clean_response_content(message.content)
             
             st.markdown(f"""
             <div class="assistant-message">
                 🤖 <strong>Assistant:</strong><br>
-                {formatted_content}
+                {content}
             </div>
             """, unsafe_allow_html=True)
     
@@ -84,16 +107,15 @@ def display_chat_interface():
                 st.session_state.messages.append(human_message)
                 
                 # Get response
-                response = st.session_state.qa_system.invoke(
-                    {
-                        "input": prompt,
-                        "chat_history": st.session_state.messages
-                    } # Missing comma here!
-                )
+                response = st.session_state.qa_system.invoke({
+                    "input": prompt,
+                    "chat_history": st.session_state.messages
+                })
                 
                 if response:
-                    response_str = str(response)
-                    ai_message = AIMessage(content=response_str)
+                    # Clean the response before creating the message
+                    cleaned_response = clean_response_content(response)
+                    ai_message = AIMessage(content=cleaned_response)
                     st.session_state.messages.append(ai_message)
                     st.rerun()
                 else:
@@ -102,4 +124,5 @@ def display_chat_interface():
         except Exception as e:
             st.error(f"Error: {e}")
             import traceback
-            st.error(traceback.format_exc())
+            st.error(traceback.format_exc())
+