Update components/chat.py

components/chat.py

@@ -2,71 +2,75 @@ import streamlit as st
 from langchain_core.messages import HumanMessage, AIMessage
 from utils.database import verify_vector_store
 
-def format_assistant_response(content):
-    """Format the assistant's response into a structured layout."""
-    # Clean up the content first
+def clean_ai_response(content):
+    """Clean up AI response content and remove technical artifacts."""
     content = str(content)
     
-    # Remove metadata and extra markers
+    # Remove common technical artifacts
     if "content='" in content:
-        content = content.split("content='")[1].split("additional_kwargs")[0].strip("'")
-    
-    # Split into sections
-    sections = {}
-    current_section = "default"
-    lines = content.replace('\\n', '\n').split('\n')
+        content = content.split("content='")[1]
+    if "additional_kwargs" in content:
+        content = content.split("additional_kwargs")[0]
     
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-            
-        if line.startswith('**') and line.endswith('**'):
-            current_section = line.strip('**')
-            sections[current_section] = []
-        elif line.startswith('*') or line.startswith('#'):
-            current_section = "main"
-            if current_section not in sections:
-                sections[current_section] = []
-            sections[current_section].append(line.strip('* '))
-        else:
-            if current_section not in sections:
-                sections[current_section] = []
-            sections[current_section].append(line)
+    # Clean up any remaining artifacts
+    content = content.strip("'")
+    content = content.replace('\\n', '\n')
+    content = content.replace('\\t', '\t')
     
-    # Build formatted HTML
-    formatted_html = """
-    <div class="response-container">
-        <div class="summary-section">
-            <h4>Executive Summary</h4>
-            <p>{}</p>
+    return content
+
+def format_assistant_response(content):
+    """Format the assistant's response into a structured layout."""
+    try:
+        # Clean the content first
+        content = clean_ai_response(content)
+        
+        # Identify sections and structure
+        lines = [line.strip() for line in content.split('\n') if line.strip()]
+        formatted_sections = []
+        current_section = []
+        
+        for line in lines:
+            # Handle bullet points and lists
+            if line.startswith(('•', '-', '*')):
+                line = f"<li>{line.lstrip('•-* ')}</li>"
+                if not current_section:
+                    current_section.append("<ul>")
+                current_section.append(line)
+            # Handle section headers
+            elif line.startswith('**') and line.endswith('**'):
+                if current_section:
+                    if current_section[0] == "<ul>":
+                        current_section.append("</ul>")
+                    formatted_sections.extend(current_section)
+                    current_section = []
+                header = line.strip('**')
+                formatted_sections.append(f'<h4 class="section-header">{header}</h4>')
+            # Regular text
+            else:
+                if current_section and current_section[0] == "<ul>":
+                    current_section.append("</ul>")
+                    formatted_sections.extend(current_section)
+                    current_section = []
+                formatted_sections.append(f'<p>{line}</p>')
+        
+        # Add any remaining content
+        if current_section:
+            if current_section[0] == "<ul>":
+                current_section.append("</ul>")
+            formatted_sections.extend(current_section)
+        
+        # Build the final HTML
+        formatted_html = f"""
+        <div class="response-content">
+            {''.join(formatted_sections)}
         </div>
-    """.format(' '.join(sections.get('Executive Summary', sections.get('default', ['']))))
-    
-    # Add key themes section
-    if 'Key Themes' in sections:
-        formatted_html += """
-        <div class="themes-section">
-            <h4>Key Themes</h4>
-            <ul>
         """
-        for theme in sections['Key Themes']:
-            if theme.strip():
-                formatted_html += f'<li class="theme-item">{theme}</li>'
-        formatted_html += '</ul></div>'
-    
-    # Add any remaining sections
-    for section, content in sections.items():
-        if section not in ['Executive Summary', 'Key Themes', 'default']:
-            formatted_html += f"""
-            <div class="content-section">
-                <h4>{section}</h4>
-                <p>{'<br>'.join(content)}</p>
-            </div>
-            """
-    
-    formatted_html += '</div>'
-    return formatted_html
+        
+        return formatted_html
+    except Exception as e:
+        st.error(f"Error formatting response: {str(e)}")
+        return str(content)  # Return raw content if formatting fails
 
 def display_chat_interface():
     """Display modern chat interface with clean formatting."""
@@ -74,103 +78,138 @@ def display_chat_interface():
     # Add custom CSS for modern chat styling
     st.markdown("""
     <style>
-        .response-container {
-            background-color: white;
+        /* Chat container */
+        .chat-container {
+            max-width: 800px;
+            margin: auto;
+        }
+        
+        /* User message */
+        .user-message {
+            background-color: #f0f2f6;
+            padding: 1rem;
             border-radius: 10px;
-            padding: 20px;
-            margin: 10px 0;
-            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+            margin: 1rem 0;
+            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
         }
         
-        .summary-section {
-            background-color: #f8f9fa;
-            padding: 15px;
-            border-radius: 8px;
-            margin-bottom: 20px;
+        /* Assistant message */
+        .assistant-message {
+            background-color: #ffffff;
+            border: 1px solid #e0e0e0;
+            padding: 1.5rem;
+            border-radius: 10px;
+            margin: 1rem 0;
+            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
         }
         
-        .themes-section {
-            margin: 20px 0;
+        /* Response content */
+        .response-content {
+            line-height: 1.6;
         }
         
-        .theme-item {
-            margin: 10px 0;
-            padding: 5px 0;
-            line-height: 1.5;
+        /* Section headers */
+        .section-header {
+            color: #0f52ba;
+            margin: 1.5rem 0 1rem 0;
+            font-size: 1.1rem;
+            font-weight: 600;
+            border-bottom: 2px solid #e0e0e0;
+            padding-bottom: 0.5rem;
         }
         
-        .content-section {
-            margin: 15px 0;
-            padding: 10px;
-            border-left: 3px solid #0f52ba;
+        /* Lists */
+        .response-content ul {
+            margin: 1rem 0;
+            padding-left: 1.5rem;
+            list-style-type: none;
         }
         
-        .content-section h4 {
+        .response-content li {
+            margin: 0.5rem 0;
+            position: relative;
+            padding-left: 1rem;
+        }
+        
+        .response-content li:before {
+            content: "•";
+            position: absolute;
+            left: -1rem;
             color: #0f52ba;
-            margin-bottom: 10px;
         }
         
-        .user-message {
-            background-color: #f0f2f6;
-            padding: 1rem;
-            border-radius: 10px;
+        /* Paragraphs */
+        .response-content p {
             margin: 1rem 0;
+            color: #2c3e50;
         }
         
-        .assistant-message {
-            background-color: #ffffff;
-            border: 1px solid #e0e0e0;
-            border-radius: 10px;
-            margin: 1rem 0;
+        /* Source citations */
+        .source-citation {
+            font-style: italic;
+            color: #666;
+            border-top: 1px solid #e0e0e0;
+            margin-top: 1rem;
+            padding-top: 0.5rem;
+            font-size: 0.9rem;
         }
     </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:
-        st.warning("Please upload documents first to initialize the chat system.")
-        return
-        
-    # Initialize chat history
-    if 'messages' not in st.session_state:
-        st.session_state.messages = []
-    
-    # Display chat history
-    for message in st.session_state.messages:
-        if isinstance(message, HumanMessage):
-            st.markdown(f"""
-            <div class="user-message">
-                🧑‍💼 <strong>You:</strong><br>{message.content}
-            </div>
-            """, unsafe_allow_html=True)
-        elif isinstance(message, AIMessage):
-            st.markdown(f"""
-            <div class="assistant-message">
-                🤖 <strong>Assistant:</strong>
-                {format_assistant_response(message.content)}
-            </div>
-            """, unsafe_allow_html=True)
-    
-    # Chat input at the bottom
-    if prompt := st.chat_input("Ask about your documents..."):
-        try:
+    try:
+        # Check if QA system is initialized
+        if 'qa_system' not in st.session_state or st.session_state.qa_system is None:
+            st.warning("Please upload documents first to initialize the chat system.")
+            return
+            
+        # Initialize chat history
+        if 'messages' not in st.session_state:
+            st.session_state.messages = []
+        
+        # Display chat history
+        for message in st.session_state.messages:
+            if isinstance(message, HumanMessage):
+                st.markdown(f"""
+                <div class="user-message">
+                    🧑‍💼 <strong>You:</strong><br>{message.content}
+                </div>
+                """, unsafe_allow_html=True)
+            elif isinstance(message, AIMessage):
+                st.markdown(f"""
+                <div class="assistant-message">
+                    🤖 <strong>Assistant:</strong>
+                    {format_assistant_response(message.content)}
+                </div>
+                """, unsafe_allow_html=True)
+        
+        # Chat input
+        if prompt := st.chat_input("Ask about your documents..."):
             with st.spinner("Analyzing..."):
+                # Validate input
+                if not prompt.strip():
+                    st.warning("Please enter a valid question.")
+                    return
+                
+                # Create and append human message
                 human_message = HumanMessage(content=prompt)
                 st.session_state.messages.append(human_message)
                 
+                # Get response from QA system
                 response = st.session_state.qa_system.invoke({
                     "input": prompt,
                     "chat_history": st.session_state.messages
                 })
                 
+                # Handle response
                 if response:
                     ai_message = AIMessage(content=str(response))
                     st.session_state.messages.append(ai_message)
                     st.rerun()
                 else:
-                    st.error("No valid response received")
+                    st.error("No valid response received. Please try again.")
                     
-        except Exception as e:
-            st.error(f"Error: {e}")
-            import traceback
-            st.error(traceback.format_exc())
+    except Exception as e:
+        st.error("An error occurred during chat processing.")
+        st.error(f"Error details: {str(e)}")
+        import traceback
+        st.error(traceback.format_exc())