Update components/chat.py

components/chat.py

@@ -2,19 +2,114 @@ 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
+    content = str(content)
+    
+    # Remove metadata and extra markers
+    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')
+    
+    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)
+    
+    # Build formatted HTML
+    formatted_html = """
+    <div class="response-container">
+        <div class="summary-section">
+            <h4>Executive Summary</h4>
+            <p>{}</p>
+        </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
+
 def display_chat_interface():
     """Display modern chat interface with clean formatting."""
     
     # Add custom CSS for modern chat styling
     st.markdown("""
     <style>
-        /* Clean, modern chat container */
-        .chat-container {
-            max-width: 800px;
-            margin: auto;
+        .response-container {
+            background-color: white;
+            border-radius: 10px;
+            padding: 20px;
+            margin: 10px 0;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }
+        
+        .summary-section {
+            background-color: #f8f9fa;
+            padding: 15px;
+            border-radius: 8px;
+            margin-bottom: 20px;
+        }
+        
+        .themes-section {
+            margin: 20px 0;
+        }
+        
+        .theme-item {
+            margin: 10px 0;
+            padding: 5px 0;
+            line-height: 1.5;
+        }
+        
+        .content-section {
+            margin: 15px 0;
+            padding: 10px;
+            border-left: 3px solid #0f52ba;
+        }
+        
+        .content-section h4 {
+            color: #0f52ba;
+            margin-bottom: 10px;
         }
         
-        /* Message styling */
         .user-message {
             background-color: #f0f2f6;
             padding: 1rem;
@@ -25,35 +120,9 @@ def display_chat_interface():
         .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)
     
@@ -75,11 +144,10 @@ def display_chat_interface():
             </div>
             """, unsafe_allow_html=True)
         elif isinstance(message, AIMessage):
-            content = str(message.content)
             st.markdown(f"""
             <div class="assistant-message">
-                🤖 <strong>Assistant:</strong><br>
-                {content}
+                🤖 <strong>Assistant:</strong>
+                {format_assistant_response(message.content)}
             </div>
             """, unsafe_allow_html=True)
     
@@ -87,19 +155,16 @@ def display_chat_interface():
     if prompt := st.chat_input("Ask about your documents..."):
         try:
             with st.spinner("Analyzing..."):
-                # Create message objects
                 human_message = HumanMessage(content=prompt)
                 st.session_state.messages.append(human_message)
                 
-                # Get response
                 response = st.session_state.qa_system.invoke({
                     "input": prompt,
                     "chat_history": st.session_state.messages
                 })
                 
                 if response:
-                    content = str(response)
-                    ai_message = AIMessage(content=content)
+                    ai_message = AIMessage(content=str(response))
                     st.session_state.messages.append(ai_message)
                     st.rerun()
                 else: