components/chat.py

# 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
from typing import Optional
import traceback
from utils.response_formatter import ResponseFormatter, display_formatted_response

# Create a lock for QA system access
qa_lock = Lock()

def display_chat_interface():
    """
    Display a modern chat interface with error handling and clean formatting.
    Manages conversation state and handles document context.
    """
    # Verify QA system initialization
    if not _verify_chat_ready():
        return

    # Initialize chat history if needed
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # Display existing chat history
    _display_chat_history()

    # Handle new user input
    if prompt := st.chat_input("Ask about your documents..."):
        _process_user_message(prompt)

def _verify_chat_ready() -> bool:
    """Check if the chat system is properly 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 False
    return True

def _process_user_message(prompt: str):
    """Process a new user message and generate AI response."""
    try:
        with st.spinner("Analyzing..."):
            human_message = HumanMessage(content=prompt)
            st.session_state.messages.append(human_message)
            with st.chat_message("user"):
                st.write(prompt)

            with qa_lock:
                response = st.session_state.qa_system.invoke({
                    "input": prompt,
                    "chat_history": st.session_state.messages
                })

            if response:
                # Format the response
                ai_message = AIMessage(content=str(response))
                st.session_state.messages.append(ai_message)
                with st.chat_message("assistant"):
                    display_formatted_response(
                        str(response),
                        metadata=getattr(response, 'metadata', None)
                    )

    except Exception as e:
        st.error(f"An error occurred while processing your message: {str(e)}")
        if st.session_state.get('debug_mode'):
            st.error(traceback.format_exc())

def _display_chat_history():
    """Display all messages in the chat history."""
    for message in st.session_state.messages:
        if isinstance(message, HumanMessage):
            with st.chat_message("user"):
                st.write(message.content)
        elif isinstance(message, AIMessage):
            with st.chat_message("assistant"):
                st.write(message.content)