Update components/chat.py

components/chat.py

@@ -1,76 +1,12 @@
+# components_chat.py
+
 import streamlit as st
 from langchain_core.messages import HumanMessage, AIMessage
 from utils.database import verify_vector_store
+from threading import Lock
 
-def clean_ai_response(content):
-    """Clean up AI response content and remove technical artifacts."""
-    content = str(content)
-    
-    # Remove common technical artifacts
-    if "content='" in content:
-        content = content.split("content='")[1]
-    if "additional_kwargs" in content:
-        content = content.split("additional_kwargs")[0]
-    
-    # Clean up any remaining artifacts
-    content = content.strip("'")
-    content = content.replace('\\n', '\n')
-    content = content.replace('\\t', '\t')
-    
-    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>
-        """
-        
-        return formatted_html
-    except Exception as e:
-        st.error(f"Error formatting response: {str(e)}")
-        return str(content)  # Return raw content if formatting fails
+# Create a lock for QA system access
+qa_lock = Lock()
 
 def display_chat_interface():
     """Display modern chat interface with clean formatting."""
@@ -78,138 +14,105 @@ def display_chat_interface():
     # Add custom CSS for modern chat styling
     st.markdown("""
     <style>
-        /* Chat container */
+        /* Clean, modern chat container */
         .chat-container {
             max-width: 800px;
             margin: auto;
         }
         
-        /* User message */
+        /* Message styling */
         .user-message {
             background-color: #f0f2f6;
             padding: 1rem;
             border-radius: 10px;
             margin: 1rem 0;
-            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
         }
         
-        /* 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);
-        }
-        
-        /* Response content */
-        .response-content {
-            line-height: 1.6;
         }
         
-        /* Section headers */
+        /* Section styling */
         .section-header {
+            font-weight: bold;
             color: #0f52ba;
-            margin: 1.5rem 0 1rem 0;
-            font-size: 1.1rem;
-            font-weight: 600;
-            border-bottom: 2px solid #e0e0e0;
-            padding-bottom: 0.5rem;
+            margin-top: 1rem;
         }
         
-        /* Lists */
-        .response-content ul {
-            margin: 1rem 0;
-            padding-left: 1.5rem;
-            list-style-type: none;
+        .source-info {
+            background-color: #f8f9fa;
+            padding: 0.5rem;
+            border-left: 3px solid #0f52ba;
+            margin-top: 1rem;
+            font-size: 0.9rem;
         }
         
-        .response-content li {
+        /* Content sections */
+        .content-section {
             margin: 0.5rem 0;
-            position: relative;
-            padding-left: 1rem;
-        }
-        
-        .response-content li:before {
-            content: "•";
-            position: absolute;
-            left: -1rem;
-            color: #0f52ba;
         }
         
-        /* Paragraphs */
-        .response-content p {
-            margin: 1rem 0;
-            color: #2c3e50;
-        }
-        
-        /* 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;
+        /* Clean bullets */
+        .clean-bullet {
+            margin: 0.3rem 0;
         }
     </style>
     """, unsafe_allow_html=True)
     
-    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..."):
+    # 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):
+            content = str(message.content)
+            st.markdown(f"""
+            <div class="assistant-message">
+                🤖 <strong>Assistant:</strong><br>
+                {content}
+            </div>
+            """, unsafe_allow_html=True)
+    
+    # Chat input at the bottom
+    if prompt := st.chat_input("Ask about your documents..."):
+        try:
             with st.spinner("Analyzing..."):
-                # Validate input
-                if not prompt.strip():
-                    st.warning("Please enter a valid question.")
-                    return
-                
-                # Create and append human message
+                # Create message objects
                 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
+
+                # Use lock to access QA system
+                with qa_lock:
+                    response = st.session_state.qa_system.invoke({
+                        "input": prompt,
+                        "chat_history": st.session_state.messages
+                    })
+
                 if response:
-                    ai_message = AIMessage(content=str(response))
+                    content = str(response)
+                    ai_message = AIMessage(content=content)
                     st.session_state.messages.append(ai_message)
-                    st.rerun()
+                    st.experimental_rerun()
                 else:
-                    st.error("No valid response received. Please try again.")
+                    st.error("No valid response received")
                     
-    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())
+        except Exception as e:
+            st.error(f"Error: {e}")
+            import traceback
+            st.error(traceback.format_exc())