app.py

#!/usr/bin/env python3
"""
QUADRANT RAG System - Enhanced UI v2 with Document Library
Professional chat interface with persistent document storage
"""

import os
import streamlit as st
import json
import uuid
import time
from typing import List, Dict, Any, Optional
from pathlib import Path
from datetime import datetime, timezone
import tempfile
import base64

# Load environment variables first
import os
from dotenv import load_dotenv
load_dotenv()

# Import RAG components
from rag_core import DynamicRAG, extract_pdf_pages, create_chunks

# Page configuration
st.set_page_config(
    page_title="QUADRANT RAG - AI Document Assistant",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Enhanced CSS for modern UI
st.markdown("""
<style>
    /* CSS Variables for theme */
    :root {
        --primary-color: #5468ff;
        --secondary-color: #6c63ff;
        --accent-color: #00d4ff;
        --background-color: #f8f9fa;
        --card-background: #ffffff;
        --text-primary: #2c3e50;
        --text-secondary: #718096;
        --border-color: #e2e8f0;
        --shadow-sm: 0 2px 4px rgba(0,0,0,0.05);
        --shadow-md: 0 4px 12px rgba(0,0,0,0.08);
        --shadow-lg: 0 10px 30px rgba(0,0,0,0.1);
    }
    
    /* Reset and base styles */
    .main {
        padding: 0;
        background-color: var(--background-color);
    }
    
    .stApp {
        background-color: var(--background-color);
    }
    
    /* Sidebar styling */
    section[data-testid="stSidebar"] {
        background-color: var(--card-background);
        border-right: 1px solid var(--border-color);
        box-shadow: 2px 0 5px rgba(0,0,0,0.05);
    }
    
    section[data-testid="stSidebar"] .block-container {
        padding: 1.5rem 1rem;
    }
    
    /* Header */
    .main-header {
        background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
        color: white;
        padding: 1.5rem 2rem;
        margin: -1rem -1rem 1rem -1rem;
        box-shadow: var(--shadow-md);
    }
    
    .main-header h1 {
        margin: 0;
        font-size: 2rem;
        font-weight: 700;
    }
    
    .main-header p {
        margin: 0.5rem 0 0 0;
        opacity: 0.9;
        font-size: 1.1rem;
    }
    
    /* Document library styles */
    .doc-library-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        margin: -1.5rem -1rem 1rem -1rem;
        padding: 1.5rem 1rem;
        color: white;
    }
    
    .doc-library-header h3 {
        margin: 0;
        font-size: 1.3rem;
        font-weight: 600;
    }
    
    .doc-count {
        font-size: 0.9rem;
        opacity: 0.9;
        margin-top: 0.25rem;
    }
    
    /* Document cards in sidebar */
    .doc-card {
        background: var(--card-background);
        border: 1px solid var(--border-color);
        border-radius: 8px;
        padding: 1rem;
        margin-bottom: 0.75rem;
        cursor: pointer;
        transition: all 0.2s ease;
    }
    
    .doc-card:hover {
        transform: translateY(-2px);
        box-shadow: var(--shadow-md);
        border-color: var(--primary-color);
    }
    
    .doc-card.active {
        border-color: var(--primary-color);
        background: linear-gradient(to right, #f0f4ff 0%, #ffffff 100%);
        box-shadow: var(--shadow-md);
    }
    
    .doc-card-title {
        font-weight: 600;
        color: var(--text-primary);
        margin-bottom: 0.25rem;
        font-size: 0.95rem;
    }
    
    .doc-card-meta {
        font-size: 0.8rem;
        color: var(--text-secondary);
        display: flex;
        justify-content: space-between;
        align-items: center;
    }
    
    .doc-card-date {
        font-size: 0.75rem;
        color: var(--text-secondary);
        margin-top: 0.25rem;
    }
    
    /* Chat interface */
    .chat-container {
        height: calc(100vh - 200px);
        display: flex;
        flex-direction: column;
        background: var(--card-background);
        border-radius: 12px;
        box-shadow: var(--shadow-lg);
        margin: 1rem;
        overflow: hidden;
    }
    
    .chat-header {
        background: linear-gradient(135deg, #f0f4ff 0%, #e8eeff 100%);
        border-bottom: 1px solid var(--border-color);
        padding: 1.5rem;
    }
    
    .chat-header-title {
        font-size: 1.1rem;
        font-weight: 600;
        color: var(--text-primary);
        margin: 0;
        white-space: nowrap;
        overflow: hidden;
        text-overflow: ellipsis;
    }
    
    .chat-header-subtitle {
        font-size: 0.9rem;
        color: var(--text-secondary);
        margin-top: 0.25rem;
    }
    
    .chat-messages {
        flex: 1;
        overflow-y: auto;
        padding: 1.5rem;
        background: #fafbfc;
    }
    
    /* Message styles */
    .message {
        margin-bottom: 1.5rem;
        animation: fadeInUp 0.3s ease;
    }
    
    @keyframes fadeInUp {
        from {
            opacity: 0;
            transform: translateY(10px);
        }
        to {
            opacity: 1;
            transform: translateY(0);
        }
    }
    
    .message.user {
        display: flex;
        justify-content: flex-end;
    }
    
    .message.assistant {
        display: flex;
        justify-content: flex-start;
    }
    
    .message-content {
        max-width: 70%;
        padding: 1rem 1.25rem;
        border-radius: 18px;
        position: relative;
        animation: scaleIn 0.2s ease;
    }
    
    @keyframes scaleIn {
        from {
            transform: scale(0.95);
        }
        to {
            transform: scale(1);
        }
    }
    
    .message.user .message-content {
        background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
        color: white;
        border-bottom-right-radius: 4px;
    }
    
    .message.assistant .message-content {
        background: white;
        border: 1px solid var(--border-color);
        color: var(--text-primary);
        border-bottom-left-radius: 4px;
    }
    
    /* Avatar */
    .message-avatar {
        width: 36px;
        height: 36px;
        border-radius: 50%;
        display: flex;
        align-items: center;
        justify-content: center;
        font-weight: 600;
        margin: 0 0.75rem;
    }
    
    .message.user .message-avatar {
        background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
        color: white;
    }
    
    .message.assistant .message-avatar {
        background: linear-gradient(135deg, #f0f4ff 0%, #e8eeff 100%);
        color: var(--primary-color);
    }
    
    /* Citations */
    .citations {
        margin-top: 0.75rem;
        padding-top: 0.75rem;
        border-top: 1px solid rgba(0,0,0,0.1);
    }
    
    .citation-item {
        background: rgba(0,0,0,0.05);
        padding: 0.5rem 0.75rem;
        border-radius: 8px;
        margin-top: 0.5rem;
        font-size: 0.85rem;
        border-left: 3px solid var(--accent-color);
    }
    
    /* Input area */
    .chat-input-container {
        border-top: 1px solid var(--border-color);
        background: white;
        padding: 1.5rem;
    }
    
    .chat-input-wrapper {
        display: flex;
        gap: 0.75rem;
        align-items: flex-end;
    }
    
    .chat-input {
        flex: 1;
        background: var(--background-color);
        border: 2px solid var(--border-color);
        border-radius: 12px;
        padding: 0.75rem 1rem;
        font-size: 1rem;
        transition: all 0.2s ease;
        resize: none;
        min-height: 50px;
    }
    
    .chat-input:focus {
        outline: none;
        border-color: var(--primary-color);
        background: white;
    }
    
    .send-button {
        background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
        color: white;
        border: none;
        border-radius: 12px;
        padding: 0.75rem 1.5rem;
        font-size: 1rem;
        font-weight: 600;
        cursor: pointer;
        transition: all 0.2s ease;
        display: flex;
        align-items: center;
        gap: 0.5rem;
    }
    
    .send-button:hover {
        transform: translateY(-2px);
        box-shadow: var(--shadow-md);
    }
    
    .send-button:active {
        transform: translateY(0);
    }
    
    /* Typing indicator */
    .typing-indicator {
        display: inline-flex;
        padding: 0.75rem 1rem;
        background: white;
        border-radius: 18px;
        border: 1px solid var(--border-color);
        gap: 4px;
    }
    
    .typing-dot {
        width: 8px;
        height: 8px;
        background: var(--primary-color);
        border-radius: 50%;
        animation: typing 1.4s infinite;
    }
    
    .typing-dot:nth-child(1) { animation-delay: -0.32s; }
    .typing-dot:nth-child(2) { animation-delay: -0.16s; }
    
    @keyframes typing {
        0%, 80%, 100% {
            opacity: 0.5;
            transform: scale(0.8);
        }
        40% {
            opacity: 1;
            transform: scale(1);
        }
    }
    
    /* Upload area */
    .upload-area {
        border: 2px dashed var(--primary-color);
        border-radius: 12px;
        padding: 3rem;
        text-align: center;
        background: linear-gradient(to bottom, #f0f4ff 0%, #ffffff 100%);
        transition: all 0.3s ease;
        margin: 1rem;
    }
    
    .upload-area:hover {
        border-color: var(--secondary-color);
        background: linear-gradient(to bottom, #e8eeff 0%, #f8f9ff 100%);
    }
    
    .upload-icon {
        font-size: 4rem;
        color: var(--primary-color);
        margin-bottom: 1rem;
    }
    
    /* Empty state */
    .empty-state {
        text-align: center;
        padding: 3rem;
        color: var(--text-secondary);
    }
    
    .empty-state-icon {
        font-size: 4rem;
        opacity: 0.3;
        margin-bottom: 1rem;
    }
    
    /* Buttons */
    .stButton > button {
        background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
        color: white;
        border: none;
        padding: 0.75rem 1.5rem;
        border-radius: 8px;
        font-weight: 600;
        transition: all 0.2s ease;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: var(--shadow-md);
    }
    
    /* Hide Streamlit defaults */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    
    /* Scrollbar */
    ::-webkit-scrollbar {
        width: 8px;
        height: 8px;
    }
    
    ::-webkit-scrollbar-track {
        background: var(--background-color);
    }
    
    ::-webkit-scrollbar-thumb {
        background: var(--border-color);
        border-radius: 4px;
    }
    
    ::-webkit-scrollbar-thumb:hover {
        background: var(--text-secondary);
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'rag_system' not in st.session_state:
    st.session_state.rag_system = None
if 'current_doc' not in st.session_state:
    st.session_state.current_doc = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'all_documents' not in st.session_state:
    st.session_state.all_documents = []
if 'processing' not in st.session_state:
    st.session_state.processing = False
if 'waiting_for_response' not in st.session_state:
    st.session_state.waiting_for_response = False

def init_rag_system():
    """Initialize the RAG system"""
    try:
        # Check environment variables
        from dotenv import load_dotenv
        load_dotenv()  # Reload environment variables

        openai_key = os.environ.get('OPENAI_API_KEY', '')
        qdrant_url = os.environ.get('QDRANT_URL', '')
        qdrant_key = os.environ.get('QDRANT_API_KEY', '')

        if not openai_key or openai_key == 'your-openai-api-key-here':
            st.error("❌ OpenAI API key not configured. Please set OPENAI_API_KEY in your environment.")
            return False

        if not qdrant_url or not qdrant_key:
            st.warning("⚠️ Qdrant Cloud credentials not found. Using local file storage.")

        # Show initialization progress
        progress_placeholder = st.empty()
        with progress_placeholder:
            with st.spinner("🔄 Initializing RAG System..."):
                try:
                    st.session_state.rag_system = DynamicRAG()
                    # Load all documents from Qdrant
                    st.session_state.all_documents = st.session_state.rag_system.get_all_documents()
                except Exception as init_error:
                    st.error(f"❌ RAG System initialization failed: {str(init_error)}")
                    # Continue anyway for basic functionality
                    st.session_state.all_documents = []

        progress_placeholder.success("✅ RAG System initialized successfully!")
        return True
    except Exception as e:
        st.error(f"❌ Failed to initialize RAG system: {str(e)}")
        # Don't fail completely - allow app to show error state
        return True

def process_pdf_upload(uploaded_file) -> Optional[Dict[str, Any]]:
    """Process uploaded PDF file"""
    try:
        st.session_state.processing = True
        
        # Save uploaded file
        temp_path = Path(tempfile.gettempdir()) / f"{uuid.uuid4().hex}.pdf"
        with open(temp_path, "wb") as f:
            f.write(uploaded_file.getvalue())
        
        # Extract text
        pages = extract_pdf_pages(str(temp_path))
        
        # Create chunks
        chunks = create_chunks(pages, chunk_size=3000, overlap=200)
        
        # Generate document ID
        doc_id = f"{uploaded_file.name.replace('.pdf', '')}_{int(time.time())}"
        
        # Store in Qdrant
        st.session_state.rag_system.store_document(doc_id, chunks)
        
        # Create document info
        doc_info = {
            'doc_id': doc_id,
            'title': uploaded_file.name,
            'pages': len(pages),
            'chunks': len(chunks),
            'upload_time': datetime.now(timezone.utc).isoformat()
        }
        
        # Update documents list
        st.session_state.all_documents = st.session_state.rag_system.get_all_documents()
        
        # Clean up
        temp_path.unlink()
        
        return doc_info
        
    except Exception as e:
        st.error(f"Error processing PDF: {str(e)}")
        return None
    finally:
        st.session_state.processing = False

def query_document(question: str) -> tuple[str, List[Dict[str, Any]]]:
    """Query the current document"""
    try:
        if not st.session_state.current_doc:
            return "Please select a document first.", []
        
        # Search in current document - increased for better coverage
        search_results = st.session_state.rag_system.search(
            query=question,
            doc_id=st.session_state.current_doc['doc_id'],
            top_k=10
        )
        
        if not search_results:
            return "I couldn't find relevant information about that in the document.", []
        
        # Generate answer
        answer = st.session_state.rag_system.generate_answer(question, search_results)
        
        # Check if the answer indicates insufficient evidence
        insufficient_keywords = ["insufficient evidence", "couldn't find", "no relevant information", "cannot answer"]
        
        # Prepare citations only if the answer has sufficient evidence
        citations = []
        if not any(keyword in answer.lower() for keyword in insufficient_keywords):
            for i, result in enumerate(search_results[:3]):
                citations.append({
                    'page': result['page'],
                    'text': result['text'][:150] + "..." if len(result['text']) > 150 else result['text'],
                    'score': round(result['score'], 3)
                })
        
        return answer, citations
        
    except Exception as e:
        return f"Sorry, I encountered an error: {str(e)}", []

def render_sidebar():
    """Render the document library sidebar"""
    with st.sidebar:
        # Header
        st.markdown("""
        <div class="doc-library-header">
            <h3>📚 Document Library</h3>
            <div class="doc-count">{} documents stored</div>
        </div>
        """.format(len(st.session_state.all_documents)), unsafe_allow_html=True)
        
        # Upload new document
        with st.expander("📤 Upload New Document", expanded=False):
            uploaded_file = st.file_uploader(
                "Choose a PDF file",
                type=['pdf'],
                label_visibility="collapsed",
                disabled=st.session_state.processing
            )
            
            if uploaded_file and st.button("Upload", type="primary", use_container_width=True):
                with st.spinner("Processing..."):
                    doc = process_pdf_upload(uploaded_file)
                    if doc:
                        st.success("✅ Document uploaded successfully!")
                        st.rerun()
        
        # Document list
        if st.session_state.all_documents:
            st.markdown("### Your Documents")
            
            for doc in st.session_state.all_documents:
                # Check if this is the current document
                is_active = (st.session_state.current_doc and 
                           doc['doc_id'] == st.session_state.current_doc['doc_id'])
                
                # Document card
                card_class = "doc-card active" if is_active else "doc-card"
                
                col1, col2 = st.columns([5, 1])
                
                with col1:
                    if st.button(
                        f"📄 **{doc['title'][:30]}{'...' if len(doc['title']) > 30 else ''}**\n\n"
                        f"📊 {doc['pages']} pages • {doc['chunks']} chunks",
                        key=f"doc_{doc['doc_id']}",
                        use_container_width=True
                    ):
                        st.session_state.current_doc = doc
                        st.session_state.chat_history = []
                        st.rerun()
                
                with col2:
                    if st.button("🗑️", key=f"del_{doc['doc_id']}", 
                               help="Delete this document"):
                        if st.session_state.rag_system.delete_document(doc['doc_id']):
                            st.session_state.all_documents = st.session_state.rag_system.get_all_documents()
                            if (st.session_state.current_doc and 
                                doc['doc_id'] == st.session_state.current_doc['doc_id']):
                                st.session_state.current_doc = None
                                st.session_state.chat_history = []
                            st.rerun()
            
        else:
            st.markdown("""
            <div class="empty-state">
                <div class="empty-state-icon">📭</div>
                <p>No documents yet</p>
                <p style="font-size: 0.85rem;">Upload your first PDF to get started</p>
            </div>
            """, unsafe_allow_html=True)

def render_chat_interface():
    """Render the main chat interface"""
    if not st.session_state.current_doc:
        # No document selected
        st.markdown("""
        <div class="upload-area">
            <div class="upload-icon">📚</div>
            <h2>Welcome to QUADRANT RAG Medical Assistant</h2>
            <p style="font-size: 1.1rem; color: #718096; margin-top: 1rem;">
                Upload medical documents or select from your library to start AI-powered medical Q&A
            </p>
            <p style="font-size: 0.95rem; color: #a0aec0; margin-top: 0.5rem;">
                ✨ Powered by OpenAI GPT-5-mini & Qdrant Cloud • Optimized for Medical Education
            </p>
        </div>
        """, unsafe_allow_html=True)
    else:
        # Chat interface
        title = st.session_state.current_doc['title']
        # Truncate overly long titles for cleaner UI
        display_title = (title[:100] + "…") if len(title) > 100 else title
        pages = st.session_state.current_doc['pages']
        chunks = st.session_state.current_doc['chunks']
        st.markdown(
            f"""
            <div class="chat-header">
                <div class="chat-header-title" title="{title}">💬 Chatting with: {display_title}</div>
                <div class="chat-header-subtitle">{pages} pages • {chunks} chunks • Ask anything about this document</div>
            </div>
            """,
            unsafe_allow_html=True,
        )

        # New chat UI using Streamlit's native components
        if not st.session_state.chat_history:
            st.info("Start a conversation about your document. Ask me to explain, summarize, or find specifics.")

        for msg in st.session_state.chat_history:
            if msg['type'] == 'user':
                with st.chat_message("user"):
                    st.markdown(msg['content'])
            else:
                with st.chat_message("assistant"):
                    st.markdown(msg['content'])
                    if msg.get('citations'):
                        with st.expander(f"📚 {len(msg['citations'])} Sources"):
                            for i, cite in enumerate(msg['citations'], 1):
                                st.markdown(f"**[{i}] Page {cite['page']}** (Relevance: {cite['score']:.3f})")
                                st.text(cite['text'][:200] + "..." if len(cite['text']) > 200 else cite['text'])
                                st.divider()

        # Chat input and immediate handling
        if prompt := st.chat_input("Ask anything about this document…"):
            st.session_state.chat_history.append({'type': 'user', 'content': prompt})
            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):
                    answer, citations = query_document(prompt)
                    st.session_state.chat_history.append({
                        'type': 'assistant',
                        'content': answer,
                        'citations': citations if citations else None
                    })
                    st.markdown(answer)
                    if citations:
                        with st.expander(f"📚 {len(citations)} Sources"):
                            for i, cite in enumerate(citations, 1):
                                st.markdown(f"**[{i}] Page {cite['page']}** (Relevance: {cite['score']:.3f})")
                                st.text(cite['text'][:200] + "..." if len(cite['text']) > 200 else cite['text'])
                                st.divider()

        # Prevent legacy UI from rendering below
        return
        

def main():
    # Configuration section for missing environment variables
    openai_key = os.environ.get('OPENAI_API_KEY', '')

    # Check if we're in Hugging Face Spaces environment
    is_hf_spaces = os.environ.get('SPACE_ID') is not None

    if not openai_key or openai_key == 'your-openai-api-key-here':
        if is_hf_spaces:
            st.error("🔑 **OpenAI API Key Required for Hugging Face Spaces**")
            st.markdown("""
            To use this app on Hugging Face Spaces:
            1. Go to your Space Settings
            2. Add a new secret named `OPENAI_API_KEY`
            3. Enter your OpenAI API key as the value
            4. Restart the Space

            You can get an API key from: https://platform.openai.com/api-keys
            """)
        else:
            st.error("🔑 **OpenAI API Key Required**")
            st.markdown("""
            Please set your OpenAI API key:
            1. Add `OPENAI_API_KEY=your-key-here` to the `.env` file, OR
            2. Set it as an environment variable in your deployment platform
            """)

            # Quick input for testing (only in local environment)
            with st.expander("💡 Quick Setup (for testing)"):
                key_input = st.text_input("Enter OpenAI API Key:", type="password")
                if st.button("Set API Key") and key_input:
                    os.environ['OPENAI_API_KEY'] = key_input
                    st.success("✅ API Key set! Initializing system...")
                    st.rerun()
        st.stop()

    # Initialize system (non-blocking for faster health check)
    if not st.session_state.rag_system:
        init_rag_system()  # This now doesn't block the app even if it fails
    
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>🤖 QUADRANT RAG - Document AI Assistant</h1>
        <p>Powered by Qdrant Vector Database & OpenAI GPT-4o-mini</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar
    render_sidebar()
    
    # Main content
    render_chat_interface()

if __name__ == "__main__":
    main()