streamlit_app.py

import streamlit as st
import os
import tempfile
from pathlib import Path
from LLM import RAGLLMSystem
from common_utils import load_yaml_config

config = load_yaml_config()

# Set page configuration
st.set_page_config(
    page_title="RAG Chat",
    page_icon="🤖",
    layout="centered"
)

# Initialize session state
if 'rag_system' not in st.session_state:
    st.session_state.rag_system = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'documents_uploaded' not in st.session_state:
    st.session_state.documents_uploaded = False

def initialize_rag_system():
    """Initialize the RAG system."""
    try:
        with st.spinner("Starting..."):
            st.session_state.rag_system = RAGLLMSystem()
        st.success("System ready!")
        return True
    except Exception as e:
        st.error(f"Error: {str(e)}")
        return False

def save_uploaded_file(uploaded_file):
    """Save uploaded file to temporary directory."""
    try:
        # Create uploads directory if it doesn't exist
        upload_dir = Path(config["data_directory"])
        upload_dir.mkdir(exist_ok=True)
        
        # Save file
        file_path = upload_dir / uploaded_file.name
        with open(file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        
        return str(file_path)
    except Exception as e:
        st.error(f"Error saving file {uploaded_file.name}: {str(e)}")
        return None

def main():
    st.title("🤖 RAG Chat")
    
    # Sidebar for controls
    with st.sidebar:
        # Initialize RAG system button
        if st.button("🚀 Start System", type="primary"):
            initialize_rag_system()
        
        # System status
        if st.session_state.rag_system:
            st.success("✅ Active")
        else:
            st.warning("⚠️ Not Started")
        
        st.markdown("---")
        
        # Upload section
        st.subheader("📄 Upload")
        uploaded_files = st.file_uploader(
            "PDF files",
            type=['pdf'],
            accept_multiple_files=True
        )
        
        if uploaded_files and st.session_state.rag_system:
            if st.button("Process", type="secondary"):
                process_documents(uploaded_files)
        
        # Clear chat button
        st.markdown("---")
        if st.button("Clear Chat"):
            st.session_state.chat_history = []
            st.rerun()

    # Main chat area
    # Display chat history
    for question, answer in st.session_state.chat_history:
        with st.chat_message("user"):
            st.write(question)
        with st.chat_message("assistant"):
            st.write(answer)
    
    # Chat input
    if st.session_state.rag_system:
        question = st.chat_input("Ask about your documents...")
        
        if question:
            # Add user message to chat
            with st.chat_message("user"):
                st.write(question)
            
            # Get response from RAG system
            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):
                    response = st.session_state.rag_system.query(question)
                st.write(response)
            
            # Add to chat history
            st.session_state.chat_history.append((question, response))
            st.rerun()
    else:
        st.info("Start the system first.")

def process_documents(uploaded_files):
    """Process uploaded documents."""
    if not uploaded_files:
        return
    
    with st.spinner("Processing..."):
        try:
            # Save uploaded files
            file_paths = []
            for uploaded_file in uploaded_files:
                file_path = save_uploaded_file(uploaded_file)
                if file_path:
                    file_paths.append(file_path)
            
            if not file_paths:
                st.error("No files saved.")
                return
            
            # Process documents with RAG system
            success = st.session_state.rag_system.add_documents(file_paths)
            
            if success:
                st.success(f"Added {len(file_paths)} documents!")
                st.session_state.documents_uploaded = True
            else:
                st.error("Processing failed.")
                
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()