fix ui issues

README.md

@@ -1,19 +1,147 @@
 ---
-title: Research Rag Chatbot
-emoji: 🚀
+title: ML Research Paper RAG Chatbot
+emoji: 🤖
 colorFrom: blue
 colorTo: purple
 sdk: docker
 app_port: 8501
 tags:
 - streamlit
+- machine-learning
+- research
+- rag
+- chatbot
 pinned: false
-short_description: Rag chatbot for research papers
+short_description: AI-powered chatbot for ML research papers
 ---
 
-# Welcome to Streamlit!
+# 📄 ML Research Paper RAG Chatbot
 
-Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
+An intelligent research assistant that helps you discover, understand, and explore Machine Learning research papers from ArXiv using Retrieval-Augmented Generation (RAG).
+
+## 🎯 What is this?
+
+This chatbot uses advanced AI to help you:
+- 🔍 **Find relevant research papers** on any ML topic
+- 📚 **Get detailed explanations** from published research
+- 💡 **Understand complex concepts** with cited sources
+- 🎓 **Stay updated** with ML research trends
+
+## ✨ Features
+
+- **Multi-LLM Support**: Choose between Anthropic Claude, Google Gemini, or Groq
+- **Smart Retrieval**: FAISS vector store with semantic search and reranking
+- **Research-Focused**: Only provides answers based on actual papers (no hallucinations)
+- **Citation-Backed**: All responses cite source papers with metadata
+- **Interactive UI**: Clean Streamlit interface with helpful guides
+
+## 🚀 Quick Start Guide
+
+### For First-Time Users
+
+1. **Start a conversation** by typing a question in the chat box
+2. **Try example queries** using the quick action buttons
+3. **Explore results** by expanding the "View Retrieved Documents" section
+4. **Adjust settings** in the sidebar for fine-tuned results
+
+### Example Queries
+
+```
+✅ Find papers on handling imbalanced datasets
+✅ What methods are used for fraud detection in ML?
+✅ Explain the attention mechanism in transformers
+✅ List recent papers about reinforcement learning
+✅ How does batch normalization improve training?
+```
+
+## 💡 Tips for Best Results
+
+### Ask Better Questions
+- ✅ **Be specific**: "fraud detection in credit cards" > "fraud"
+- ✅ **Use ML terminology**: "convolutional neural networks" > "image AI"
+- ✅ **Ask for comparisons**: "Compare CNN vs RNN for sequences"
+
+### Understand the Responses
+- 📚 All answers are based on research papers in the database
+- 🔍 Check "View Retrieved Documents" to see sources
+- ⚠️ If documents seem irrelevant, try rephrasing
+
+### Advanced Usage
+- ⚙️ Adjust retrieval settings (base_k, rerank_k) for more/fewer papers
+- 🎨 Switch LLM providers for different response styles
+- 📅 Filter by year or category for focused results
+
+## 🗂️ Dataset
+
+Uses **CShorten/ML-ArXiv-Papers** from Hugging Face:
+- Curated Machine Learning research papers from ArXiv
+- Includes titles, abstracts, metadata, and citations
+- Regularly updated with new publications
+
+## ⚙️ Configuration
+
+### LLM Providers
+1. **Anthropic Claude** (Recommended for quality)
+   - claude-3-5-sonnet-20241022 (Best balance)
+   - claude-3-5-haiku-20241022 (Fast)
+   
+2. **Google Gemini** (Good for free tier)
+   - gemini-2.5-flash (Fast and efficient)
+   
+3. **Groq** (Fastest inference)
+   - llama-4-maverick-17b (Open source)
+
+### Retrieval Settings
+- **base_k**: Initial papers fetched (4-30, default: 20)
+- **rerank_k**: Final papers after reranking (1-12, default: 8)
+- **Dynamic k**: Auto-adjust based on query
+- **Reranking**: Improve relevance with cross-encoder
+
+## 🔧 Setup (For Developers)
+
+### Prerequisites
+```bash
+pip install -r requirements.txt
+```
+
+### API Keys
+Create a `.env` file:
+```env
+ANTHROPIC_API_KEY=your-key-here
+GEMINI_API_KEY=your-key-here
+GROQ_API_KEY=your-key-here
+```
+
+### Run Locally
+```bash
+streamlit run streamlit_app.py
+```
+
+## 📊 How It Works
+
+1. **User Query** → Semantic embedding created
+2. **Vector Search** → FAISS retrieves similar papers
+3. **Reranking** → Cross-encoder scores relevance
+4. **LLM Generation** → AI generates answer from papers
+5. **Response** → Cited answer with source papers
+
+## 🔒 Important Notes
+
+- ✅ Answers are **based only on research papers** in the database
+- ✅ System won't make up information from general knowledge
+- ✅ If no relevant papers found, it will tell you
+- ❌ Not a replacement for reading the full papers
+- ⚠️ Always verify critical information with original sources
+
+## 🤝 Contributing
+
+Feel free to submit issues, fork the repository, and create pull requests for any improvements.
+
+## 📄 License
+
+This project is open source and available under the MIT License.
+
+---
+
+**Ready to explore ML research?** Start by asking a question! 🚀
 
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).

faiss_index/index.faiss

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5990c6b4f15d524dae50f9d256cd3e16feef0863d5bb3f467629b4b7534cdca5
-size 151790637

faiss_index/index.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d5fd0b4d6d2b61c5f67dc2447721091df6eee5a6a47b8d7018a7775b880aabd0
-size 63795991

requirements.txt

@@ -29,6 +29,7 @@ langchain-text-splitters==0.3.11
 langchain-groq==0.3.8
 langchain-huggingface==0.3.1
 langchain-google-genai==2.1.12
+langchain-anthropic==0.3.6
 
 # Vector store and NLP
 faiss-cpu==1.12.0
@@ -39,11 +40,12 @@ transformers==4.56.1
 pandas==2.3.2
 numpy==1.26.4
 requests==2.32.5
+datasets==3.2.0
 
 # Optional semantic splitter (app gracefully falls back if missing)
 semantic-text-splitter==0.27.0
 
-# Dataset fetcher
-kagglehub==0.3.13
+# Dataset fetcher (legacy - now using Hugging Face datasets)
+# kagglehub==0.3.13
 
 

src/config.py

@@ -12,10 +12,12 @@ DEVICE = "cuda" if os.environ.get("CUDA_AVAILABLE", "0") == "1" else "cpu"
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
 GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", GOOGLE_API_KEY)
+ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY", "")
 
 # Default chat model identifiers
 GROQ_MODEL = os.environ.get("GROQ_MODEL", "meta-llama/llama-4-maverick-17b-128e-instruct")
 GEMINI_MODEL = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash")
+ANTHROPIC_MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")
 
 # Cross-encoder model for reranking
 CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-12-v2"

src/ingestion.py

@@ -8,6 +8,38 @@ from langchain_core.documents import Document
 from .config import DATA_PATH
 from .text_processing import clean_text
 
+def load_hf_dataset(num_records=50000, dataset_name="CShorten/ML-ArXiv-Papers"):
+    """Load ArXiv papers from Hugging Face dataset.
+    
+    Args:
+        num_records: Number of records to load
+        dataset_name: Hugging Face dataset identifier
+    
+    Returns:
+        pandas DataFrame with the papers
+    """
+    try:
+        from datasets import load_dataset
+        
+        print(f"Loading {num_records} records from {dataset_name}...")
+        
+        # Load dataset from Hugging Face
+        dataset = load_dataset(dataset_name, split="train", streaming=False)
+        
+        # Convert to pandas DataFrame
+        if num_records and num_records < len(dataset):
+            df = dataset.select(range(num_records)).to_pandas()
+        else:
+            df = dataset.to_pandas()
+        
+        print(f"Loaded {len(df)} records from Hugging Face dataset")
+        return df
+        
+    except ImportError:
+        raise ImportError("Please install the datasets library: pip install datasets")
+    except Exception as e:
+        raise ValueError(f"Failed to load Hugging Face dataset: {e}")
+
 def _open_file(file_path):
     """Open file with appropriate mode and encoding."""
     if file_path.endswith('.gz'):
@@ -84,10 +116,28 @@ def load_data_subset(file_path, num_records=50000):
     return pd.DataFrame(records)
 
 def preprocess_dataframe(df: pd.DataFrame) -> pd.DataFrame:
-    df['update_date'] = pd.to_datetime(df['update_date'])
-    df['year'] = df['update_date'].dt.year
-    df = df.dropna(subset=['abstract'])
-    df = df[df['abstract'].str.strip() != '']
+    """Preprocess the dataframe from Hugging Face or local file."""
+    # Handle different date column names
+    date_col = None
+    if 'update_date' in df.columns:
+        date_col = 'update_date'
+    elif 'updated' in df.columns:
+        date_col = 'updated'
+    elif 'published' in df.columns:
+        date_col = 'published'
+    
+    if date_col:
+        df[date_col] = pd.to_datetime(df[date_col], errors='coerce')
+        df['year'] = df[date_col].dt.year
+    elif 'year' not in df.columns:
+        # If no date column exists, set year to None
+        df['year'] = None
+    
+    # Ensure required columns exist
+    if 'abstract' in df.columns:
+        df = df.dropna(subset=['abstract'])
+        df = df[df['abstract'].str.strip() != '']
+    
     return df
 
 def df_to_documents(
@@ -95,19 +145,35 @@ def df_to_documents(
     lowercase: bool = False,
     remove_stopwords: bool = False
 ):
+    """Convert dataframe to LangChain documents."""
     documents = []
     for _, row in df.iterrows():
-        title_clean = clean_text(str(row['title']), lowercase=lowercase, remove_stopwords=remove_stopwords)
-        abstract_clean = clean_text(str(row['abstract']), lowercase=lowercase, remove_stopwords=remove_stopwords)
+        # Get title and abstract
+        title = str(row.get('title', ''))
+        abstract = str(row.get('abstract', ''))
+        
+        title_clean = clean_text(title, lowercase=lowercase, remove_stopwords=remove_stopwords)
+        abstract_clean = clean_text(abstract, lowercase=lowercase, remove_stopwords=remove_stopwords)
         page_content = f"Title: {title_clean}\n\nAbstract: {abstract_clean}"
+        
+        # Handle categories - can be string or list
         categories_raw = row.get('categories', 'N/A') or 'N/A'
-        primary_category = categories_raw.split()[0] if isinstance(categories_raw, str) else 'N/A'
+        if isinstance(categories_raw, list):
+            categories_str = ' '.join(categories_raw) if categories_raw else 'N/A'
+            primary_category = categories_raw[0] if categories_raw else 'N/A'
+        else:
+            categories_str = str(categories_raw)
+            primary_category = categories_str.split()[0] if categories_str != 'N/A' else 'N/A'
+        
+        # Build metadata
         metadata = {
             "id": row.get('id', 'N/A'),
+            "title": title,  # Keep original title in metadata
             "authors": row.get('authors', 'N/A'),
             "year": int(row.get('year')) if not pd.isna(row.get('year')) else None,
-            "categories": categories_raw,
+            "categories": categories_str,
             "primary_category": primary_category
         }
+        
         documents.append(Document(page_content=page_content, metadata=metadata))
     return documents

streamlit_app.py

@@ -30,159 +30,540 @@ from langchain.prompts import PromptTemplate
 from langchain.schema.runnable import RunnablePassthrough
 from langchain_core.runnables import RunnableLambda
 from langchain_groq import ChatGroq
+import time
 
 from src.vector_store import build_or_load_vectorstore
-from src.ingestion import load_data_subset, preprocess_dataframe, df_to_documents
+from src.ingestion import load_data_subset, preprocess_dataframe, df_to_documents, load_hf_dataset
 from src.retriever import build_advanced_retriever
-from src.config import DATA_PATH, FAISS_INDEX_PATH, GROQ_API_KEY, GEMINI_API_KEY, GROQ_MODEL, GEMINI_MODEL
+from src.config import DATA_PATH, FAISS_INDEX_PATH, GROQ_API_KEY, GEMINI_API_KEY, ANTHROPIC_API_KEY, GROQ_MODEL, GEMINI_MODEL, ANTHROPIC_MODEL
 
 load_dotenv(find_dotenv())
 
-st.set_page_config(page_title="📄 Research Paper RAG Chatbot", page_icon="💬", layout="wide")
-st.title("📄 Research Paper RAG Chatbot (Groq + FAISS + Rerank)")
+# PAGE CONFIG - Must be first Streamlit command
+st.set_page_config(
+    page_title="Research Assistant", 
+    page_icon="🤖", 
+    layout="wide",
+    initial_sidebar_state="expanded"  # Start with sidebar expanded
+)
 
-# Sidebar controls
+# ENHANCED CUSTOM CSS - ChatGPT-like styling
+st.markdown("""
+<style>
+    /* Hide Streamlit branding */
+    #MainMenu {visibility: hidden;}
+    footer {visibility: hidden;}
+    
+    /* Make sure header is visible for sidebar toggle */
+    header {visibility: visible !important;}
+    
+    /* Style the sidebar toggle button to be more visible */
+    [data-testid="collapsedControl"] {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+        border-radius: 0 8px 8px 0 !important;
+        padding: 8px !important;
+        margin-top: 60px !important;
+        box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4) !important;
+    }
+    
+    [data-testid="collapsedControl"]:hover {
+        box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
+        transform: translateX(2px);
+    }
+    
+    /* Overall app styling */
+    .stApp {
+        background: linear-gradient(180deg, #0f1419 0%, #1a1f2e 100%);
+    }
+    
+    /* Main chat container */
+    .main .block-container {
+        padding-top: 2rem;
+        padding-bottom: 2rem;
+        max-width: 900px;
+        margin: 0 auto;
+    }
+    
+    /* Chat input styling - Fixed at bottom like ChatGPT */
+    .stChatInputContainer {
+        background: transparent;
+        border: none;
+        padding: 1rem 0;
+    }
+    
+    .stChatInput > div {
+        background: rgba(255, 255, 255, 0.05);
+        border: 1px solid rgba(255, 255, 255, 0.1);
+        border-radius: 24px;
+        padding: 12px 20px;
+        backdrop-filter: blur(10px);
+        transition: all 0.3s ease;
+    }
+    
+    .stChatInput > div:hover {
+        background: rgba(255, 255, 255, 0.08);
+        border-color: rgba(255, 255, 255, 0.2);
+        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3);
+    }
+    
+    .stChatInput > div:focus-within {
+        background: rgba(255, 255, 255, 0.1);
+        border-color: #10a37f;
+        box-shadow: 0 0 0 3px rgba(16, 163, 127, 0.1);
+    }
+    
+    /* User messages - Right aligned with gradient */
+    [data-testid="stChatMessage"]:has([data-testid*="user"]) {
+        background: transparent;
+        justify-content: flex-end;
+    }
+    
+    [data-testid="stChatMessage"]:has([data-testid*="user"]) [data-testid="stChatMessageContent"] {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border-radius: 18px;
+        padding: 14px 18px;
+        margin-left: auto;
+        max-width: 75%;
+        color: white;
+        box-shadow: 0 4px 12px rgba(102, 126, 234, 0.3);
+    }
+    
+    /* Bot messages - Left aligned with subtle styling */
+    [data-testid="stChatMessage"]:not(:has([data-testid*="user"])) {
+        background: transparent;
+        justify-content: flex-start;
+    }
+    
+    [data-testid="stChatMessage"]:not(:has([data-testid*="user"])) [data-testid="stChatMessageContent"] {
+        background: rgba(255, 255, 255, 0.03);
+        border: 1px solid rgba(255, 255, 255, 0.08);
+        border-radius: 18px;
+        padding: 14px 18px;
+        margin-right: auto;
+        max-width: 85%;
+        color: #e8e8e8;
+        box-shadow: 0 2px 8px rgba(0, 0, 0, 0.2);
+        backdrop-filter: blur(10px);
+    }
+    
+    /* Avatar styling */
+    [data-testid="stChatMessage"] [data-testid="stAvatar"] {
+        width: 36px;
+        height: 36px;
+        border-radius: 50%;
+        border: 2px solid rgba(255, 255, 255, 0.1);
+    }
+    
+    /* User avatar - gradient border */
+    [data-testid="stChatMessage"]:has([data-testid*="user"]) [data-testid="stAvatar"] {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border: 2px solid transparent;
+        box-shadow: 0 2px 8px rgba(102, 126, 234, 0.4);
+    }
+    
+    /* Bot avatar - themed */
+    [data-testid="stChatMessage"]:not(:has([data-testid*="user"])) [data-testid="stAvatar"] {
+        background: linear-gradient(135deg, #10a37f 0%, #0d8a6a 100%);
+        border: 2px solid rgba(16, 163, 127, 0.3);
+        box-shadow: 0 2px 8px rgba(16, 163, 127, 0.3);
+    }
+    
+    /* Sidebar styling */
+    [data-testid="stSidebar"] {
+        background: rgba(15, 20, 25, 0.95);
+        border-right: 1px solid rgba(255, 255, 255, 0.08);
+        backdrop-filter: blur(20px);
+    }
+    
+    [data-testid="stSidebar"] .stButton button {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border: none;
+        border-radius: 12px;
+        color: white;
+        padding: 10px 20px;
+        font-weight: 600;
+        transition: all 0.3s ease;
+    }
+    
+    [data-testid="stSidebar"] .stButton button:hover {
+        transform: translateY(-2px);
+        box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4);
+    }
+    
+    /* Expander styling */
+    .streamlit-expanderHeader {
+        background: rgba(255, 255, 255, 0.03);
+        border-radius: 12px;
+        border: 1px solid rgba(255, 255, 255, 0.08);
+        color: #b4b4b4;
+        padding: 12px 16px;
+        transition: all 0.3s ease;
+    }
+    
+    .streamlit-expanderHeader:hover {
+        background: rgba(255, 255, 255, 0.06);
+        border-color: rgba(255, 255, 255, 0.15);
+    }
+    
+    .streamlit-expanderContent {
+        background: rgba(255, 255, 255, 0.02);
+        border: 1px solid rgba(255, 255, 255, 0.05);
+        border-top: none;
+        border-radius: 0 0 12px 12px;
+    }
+    
+    /* Divider styling */
+    hr {
+        border: none;
+        height: 1px;
+        background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.1), transparent);
+        margin: 2rem 0;
+    }
+    
+    /* Info boxes */
+    .stAlert {
+        background: rgba(16, 163, 127, 0.1);
+        border: 1px solid rgba(16, 163, 127, 0.3);
+        border-radius: 12px;
+        color: #a8e6cf;
+    }
+    
+    /* Scrollbar styling */
+    ::-webkit-scrollbar {
+        width: 8px;
+        height: 8px;
+    }
+    
+    ::-webkit-scrollbar-track {
+        background: rgba(255, 255, 255, 0.02);
+    }
+    
+    ::-webkit-scrollbar-thumb {
+        background: rgba(255, 255, 255, 0.15);
+        border-radius: 10px;
+    }
+    
+    ::-webkit-scrollbar-thumb:hover {
+        background: rgba(255, 255, 255, 0.25);
+    }
+    
+    /* Typography improvements */
+    h1, h2, h3 {
+        color: #f0f0f0;
+        font-weight: 600;
+    }
+    
+    p {
+        line-height: 1.7;
+        color: #d4d4d4;
+    }
+    
+    /* Slider styling */
+    .stSlider > div > div > div > div {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    }
+    
+    /* Checkbox styling */
+    .stCheckbox > label > div[data-testid="stMarkdownContainer"] > p {
+        color: #d4d4d4;
+    }
+    
+    /* Thinking animation */
+    @keyframes pulse {
+        0%, 100% { opacity: 0.6; }
+        50% { opacity: 1; }
+    }
+    
+    .thinking {
+        animation: pulse 1.5s ease-in-out infinite;
+        color: #10a37f;
+        font-style: italic;
+    }
+    
+    /* Welcome message styling */
+    .welcome-message {
+        background: linear-gradient(135deg, rgba(16, 163, 127, 0.1) 0%, rgba(102, 126, 234, 0.1) 100%);
+        border: 1px solid rgba(16, 163, 127, 0.3);
+        border-radius: 16px;
+        padding: 24px;
+        margin: 20px 0;
+        text-align: center;
+        box-shadow: 0 4px 16px rgba(16, 163, 127, 0.1);
+    }
+    
+    .welcome-message h2 {
+        background: linear-gradient(135deg, #10a37f 0%, #667eea 100%);
+        -webkit-background-clip: text;
+        -webkit-text-fill-color: transparent;
+        margin-bottom: 12px;
+    }
+    
+    /* Suggestion chips */
+    .suggestion-chip {
+        display: inline-block;
+        background: rgba(255, 255, 255, 0.05);
+        border: 1px solid rgba(255, 255, 255, 0.1);
+        border-radius: 20px;
+        padding: 8px 16px;
+        margin: 6px;
+        color: #b4b4b4;
+        cursor: pointer;
+        transition: all 0.3s ease;
+    }
+    
+    .suggestion-chip:hover {
+        background: rgba(16, 163, 127, 0.15);
+        border-color: rgba(16, 163, 127, 0.4);
+        color: #10a37f;
+        transform: translateY(-2px);
+    }
+</style>
+""", unsafe_allow_html=True)
+
+# Title with emoji and clean design
+col1, col2, col3 = st.columns([1, 6, 1])
+with col2:
+    st.markdown("<h1 style='text-align: center; margin-bottom: 0;'>🤖 Research Assistant</h1>", unsafe_allow_html=True)
+    st.markdown("<p style='text-align: center; color: #888; margin-top: 0;'>Powered by Multi-LLM RAG + FAISS</p>", unsafe_allow_html=True)
+
+# Sidebar controls with improved organization
 with st.sidebar:
-    st.header("Retrieval Settings")
-    base_k = st.slider("Initial fetch (base_k)", 4, 30, 20, 1)  # Increased default
-    rerank_k = st.slider("Final docs (rerank_k)", 1, 12, 8, 1)  # Increased default
-    dynamic = st.checkbox("Dynamic k", True)
-    use_rerank = st.checkbox("Use reranking", True)
-    primary_category = st.text_input("Primary category filter", "") or None
-    year_min = st.number_input("Min year", value=0, step=1)
-    year_max = st.number_input("Max year", value=0, step=1)
-    if year_min == 0:
-        year_min = None
-    if year_max == 0:
-        year_max = None
-    rebuild = st.button("Rebuild index (semantic)")
-    subset_size = st.number_input("Subset records (rebuild)", 1000, 100000, 50000, 1000)
-    st.divider()
-    st.header("LLM Provider")
-    if GEMINI_API_KEY:
-        default_provider = "Gemini"
-    elif GROQ_API_KEY:
-        default_provider = "Groq"
-    else:
-        default_provider = "Gemini"
-    provider = st.selectbox("Choose provider", ["Gemini", "Groq"], index=["Gemini", "Groq"].index(default_provider))
-    if provider == "Gemini":
-        ui_gemini_model = st.text_input("Gemini model", GEMINI_MODEL)
-        ui_groq_model = None
-    else:
-        ui_groq_model = st.text_input("Groq model", GROQ_MODEL)
-        ui_gemini_model = None
+    st.markdown("### ⚙️ Configuration")
+    
+    with st.expander("📊 Dataset Info", expanded=False):
+        st.markdown("""
+        **Source:** CShorten/ML-ArXiv-Papers  
+        **Focus:** Machine Learning Research  
+        **Platform:** Hugging Face
+        """)
+    
+    st.markdown("---")
+    
+    with st.expander("🔍 Retrieval Settings", expanded=False):
+        base_k = st.slider("Initial fetch", 4, 30, 20, 1, help="Number of documents to initially retrieve")
+        rerank_k = st.slider("Final docs", 1, 12, 8, 1, help="Number of documents after reranking")
+        dynamic = st.checkbox("Dynamic k", True, help="Adjust retrieval size dynamically")
+        use_rerank = st.checkbox("Use reranking", True, help="Apply reranking for better relevance")
+        
+        with st.expander("🔧 Advanced Filters"):
+            primary_category = st.text_input("Category filter", "", help="Filter by arXiv category") or None
+            col1, col2 = st.columns(2)
+            with col1:
+                year_min = st.number_input("Min year", value=0, step=1)
+            with col2:
+                year_max = st.number_input("Max year", value=0, step=1)
+            if year_min == 0:
+                year_min = None
+            if year_max == 0:
+                year_max = None
+    
+    st.markdown("---")
+    
+    with st.expander("🔄 Index Management", expanded=False):
+        subset_size = st.number_input("Dataset size", 1000, 100000, 10000, 1000)
+        rebuild = st.button("🔨 Rebuild Index", use_container_width=True)
+    
+    st.markdown("---")
+    
+    with st.expander("🤖 LLM Provider", expanded=False):
+        # Determine default provider based on available API keys
+        if ANTHROPIC_API_KEY:
+            default_provider = "Anthropic (Claude)"
+        elif GEMINI_API_KEY:
+            default_provider = "Gemini"
+        elif GROQ_API_KEY:
+            default_provider = "Groq"
+        else:
+            default_provider = "Gemini"
+        
+        available_providers = ["Anthropic (Claude)", "Gemini", "Groq"]
+        try:
+            default_index = available_providers.index(default_provider)
+        except ValueError:
+            default_index = 0
+        
+        provider = st.selectbox("Provider", available_providers, index=default_index)
+        
+        if provider == "Anthropic (Claude)":
+            ui_anthropic_model = st.selectbox(
+                "Model",
+                [
+                    "claude-sonnet-4-5-20250929",
+                    "claude-opus-4-1-20250805",
+                    "claude-opus-4-20250514",
+                    "claude-sonnet-4-20250514",
+                    "claude-3-7-sonnet-20250219",
+                    "claude-3-5-haiku-20241022",
+                    "claude-3-haiku-20240307"
+                ],
+                index=3
+            )
+            ui_gemini_model = None
+            ui_groq_model = None
+        elif provider == "Gemini":
+            ui_gemini_model = st.text_input("Model", GEMINI_MODEL)
+            ui_groq_model = None
+            ui_anthropic_model = None
+        else:
+            ui_groq_model = st.text_input("Model", GROQ_MODEL)
+            ui_gemini_model = None
+            ui_anthropic_model = None
+    
+    # Stats at bottom
+    st.markdown("---")
+    try:
+        if 'vectorstore' in locals():
+            index_stats = vectorstore.index.ntotal if hasattr(vectorstore, 'index') else "Unknown"
+            st.metric("📚 Embeddings", f"{index_stats:,}" if isinstance(index_stats, int) else index_stats)
+    except:
+        pass
 
 # Build or load vectorstore
-def _load_df_with_fallback(data_file: str, num_records: int):
-    """Try to load the dataset; if it fails, download via KaggleHub and retry once."""
+def _load_df_from_hf(num_records: int):
+    """Load dataset from Hugging Face."""
     try:
-        return preprocess_dataframe(load_data_subset(data_file, num_records=num_records))
+        with st.spinner("🔄 Loading ML papers from Hugging Face..."):
+            df = load_hf_dataset(num_records=num_records, dataset_name="CShorten/ML-ArXiv-Papers")
+            return preprocess_dataframe(df)
     except Exception as e:
-        st.warning(f"Dataset read failed: {e}. Attempting fresh download via KaggleHub...")
-        try:
-            import kagglehub, shutil
-            os.makedirs(DATA_PATH, exist_ok=True)
-            with st.spinner("Downloading ArXiv dataset..."):
-                path = kagglehub.dataset_download("Cornell-University/arxiv")
-                src = os.path.join(path, "arxiv-metadata-oai-snapshot.json")
-                shutil.copy(src, data_file)
-            st.success("Dataset downloaded; retrying load...")
-            return preprocess_dataframe(load_data_subset(data_file, num_records=num_records))
-        except Exception as e2:
-            st.error(f"Dataset download or reload failed: {e2}")
-            st.stop()
+        st.error(f"❌ Failed to load dataset: {e}")
+        st.info("💡 Make sure 'datasets' is installed: `pip install datasets`")
+        st.stop()
 
 if rebuild or not os.path.exists(FAISS_INDEX_PATH):
-    data_file = os.path.join(DATA_PATH, "arxiv-metadata-oai-snapshot.json")
-    if not os.path.exists(data_file):
-        st.warning("Dataset missing. Attempting to download via KaggleHub...")
-        try:
-            import kagglehub, shutil
-            os.makedirs(DATA_PATH, exist_ok=True)
-            with st.spinner("Downloading ArXiv dataset..."):
-                path = kagglehub.dataset_download("Cornell-University/arxiv")
-                src = os.path.join(path, "arxiv-metadata-oai-snapshot.json")
-                shutil.copy(src, data_file)
-                st.success("Dataset downloaded.")
-        except Exception as e:
-            st.error(f"Dataset download failed: {e}. Please run main pipeline first.")
-            st.stop()
-    with st.spinner("Building vector index..."):
-        df = _load_df_with_fallback(data_file, num_records=int(subset_size))
+    with st.spinner("🔨 Building vector index..."):
+        df = _load_df_from_hf(num_records=int(subset_size))
         docs = df_to_documents(df)
         vectorstore = build_or_load_vectorstore(
             docs,
             force_rebuild=True,
             chunk_method="semantic",
-            chunk_size=800,
-            chunk_overlap=120
+            chunk_size=1000,
+            chunk_overlap=125
         )
 else:
     try:
         vectorstore = build_or_load_vectorstore([], force_rebuild=False)
     except Exception as e:
-        st.warning(f"Failed to load existing index: {e}. Attempting to rebuild from dataset...")
-        data_file = os.path.join(DATA_PATH, "arxiv-metadata-oai-snapshot.json")
-        if not os.path.exists(data_file):
-            st.error("Dataset missing. Run main pipeline first or click 'Rebuild index'.")
-            st.stop()
-        with st.spinner("Rebuilding vector index after load failure..."):
-            df = _load_df_with_fallback(data_file, num_records=50000)
+        st.warning(f"⚠️ Index load failed. Rebuilding...")
+        with st.spinner("🔨 Rebuilding vector index..."):
+            df = _load_df_from_hf(num_records=50000)
             docs = df_to_documents(df)
             vectorstore = build_or_load_vectorstore(
                 docs,
                 force_rebuild=True,
                 chunk_method="semantic",
-                chunk_size=800,
-                chunk_overlap=120
+                chunk_size=1000,
+                chunk_overlap=125
             )
 
 def make_llm(provider_name: str):
+    if provider_name == "Anthropic (Claude)":
+        if not ANTHROPIC_API_KEY:
+            st.error("❌ ANTHROPIC_API_KEY not set")
+            st.stop()
+        try:
+            from langchain_anthropic import ChatAnthropic
+            return ChatAnthropic(
+                model=ui_anthropic_model or ANTHROPIC_MODEL,
+                temperature=0.7,
+                max_tokens=2048,
+                api_key=ANTHROPIC_API_KEY,
+            )
+        except Exception as e:
+            st.error(f"❌ Claude initialization failed: {e}")
+            st.stop()
+    
     if provider_name == "Gemini":
         if not GEMINI_API_KEY:
-            st.error("GEMINI_API_KEY not set. Please add it to your .env or environment variables to use Gemini.")
+            st.error("❌ GEMINI_API_KEY not set")
             st.stop()
         try:
             from langchain_google_genai import ChatGoogleGenerativeAI
             return ChatGoogleGenerativeAI(
                 model=ui_gemini_model or GEMINI_MODEL,
                 temperature=0.7,
-                max_output_tokens=1024,  # Increased for better responses
+                max_output_tokens=1024,
                 api_key=GEMINI_API_KEY,
             )
-        except ModuleNotFoundError:
-            st.error("Missing dependency 'langchain-google-genai'. Please install it (pip install langchain-google-genai).")
-            st.stop()
         except Exception as e:
-            st.error(f"Failed to initialize Gemini: {e}")
+            st.error(f"❌ Gemini initialization failed: {e}")
             st.stop()
+    
     if not GROQ_API_KEY:
-        st.error("No valid LLM provider configured. Please set GEMINI_API_KEY or GROQ_API_KEY in environment.")
+        st.error("❌ No valid LLM provider configured")
         st.stop()
     return ChatGroq(
         model=ui_groq_model or GROQ_MODEL,
         temperature=0.7,
-        max_tokens=1024,  # Increased for better responses
+        max_tokens=1024,
         groq_api_key=GROQ_API_KEY,
     )
 
 llm = make_llm(provider)
 
-# IMPROVED PROMPT - More flexible and informative
-prompt_template = """You are a helpful research assistant specializing in arXiv papers. Your goal is to provide informative, accurate answers based on the retrieved research papers.
+# Relevance checking prompt
+relevance_check_prompt = """You are a research paper relevance checker. Your task is to determine if the retrieved documents are relevant to the user's question.
+
+Retrieved Documents:
+{context}
+
+User Question: {question}
+
+Instructions:
+- Carefully analyze whether the retrieved documents contain information that can answer the user's question
+- Consider if the documents discuss the topic, concepts, or methods mentioned in the question
+- Respond with ONLY one word: "RELEVANT" or "IRRELEVANT"
+- Be strict: if the documents are only tangentially related or don't actually address the question, respond "IRRELEVANT"
+
+Response:"""
+
+relevance_prompt = PromptTemplate(template=relevance_check_prompt, input_variables=["context", "question"])
+
+# IMPROVED PROMPT
+prompt_template = """You are a knowledgeable and helpful research assistant specializing in arXiv papers. You MUST ONLY answer questions based on the provided research papers context.
 
 Context from Research Papers:
 {context}
 
 User Question: {question}
 
+CRITICAL RULES:
+- ONLY use information from the provided research papers context above
+- DO NOT use your general knowledge or training data
+- If the context doesn't contain relevant information, you MUST respond with: "I couldn't find relevant information about this topic in the available research papers. The retrieved documents don't address your question. Please try different search terms or the database may not contain papers on this specific topic."
+
 Instructions:
-1. If the context contains relevant information, provide a comprehensive answer citing the papers (e.g., "According to Doc 1...")
-2. Mention key details like authors, year, methods, and findings when available
-3. If the context is partially relevant, provide what you can and note what's missing
-4. Structure your response clearly with bullet points or sections when appropriate
-5. Always be helpful and informative
+- Analyze the user's question and provide a thorough, well-structured response BASED ONLY ON THE CONTEXT
+- Be conversational and descriptive - explain concepts clearly with sufficient detail
+- Use multiple paragraphs when needed to fully address the question
+
+**For paper listing requests** (e.g., "find papers", "list papers", "show papers"):
+Format as a structured list with detailed summaries:
+   
+   **Paper #[Number]: [Title]**
+   - **Authors:** [Author names]
+   - **Year:** [Publication year]
+   - **ArXiv ID:** [ID if available]
+   - **Category:** [Research category]
+   - **Summary:** [3-4 sentences explaining the paper's objectives, methodology, key contributions, and findings based on the context]
+
+**For specific questions** (e.g., "What is...", "Explain...", "How does...", "What is the purpose of..."):
+- Provide a comprehensive, multi-paragraph answer that fully addresses the question USING ONLY THE CONTEXT
+- Start with a clear overview or direct answer from the papers
+- Elaborate with details, context, and explanations from the research papers
+- Discuss relevant methodologies, findings, implications, or technical details found in the papers
+- Cite sources naturally throughout (e.g., "According to the research by [Authors] (Year)...")
+- Use clear transitions between ideas
+- Conclude with key takeaways or significance when appropriate
+
+**General Guidelines:**
+- Write in a natural, conversational tone similar to ChatGPT
+- Aim for depth and clarity - don't give one-liner responses
+- Break complex information into digestible paragraphs
+- Use examples and analogies when helpful from the context
+- NEVER invent or hallucinate information not in the context
+- Always prioritize being helpful, informative, and thorough - but ONLY based on the provided context
 
 Answer:"""
 
@@ -194,35 +575,34 @@ def _format_metadata(metadata):
         return ""
     meta_lines = []
     if metadata.get("title"):
-        meta_lines.append(f"Title: {metadata['title']}")
+        meta_lines.append(f"📄 {metadata['title']}")
     if metadata.get("id"):
-        meta_lines.append(f"ArXiv ID: {metadata['id']}")
+        meta_lines.append(f"🔗 {metadata['id']}")
     if metadata.get("authors") and metadata["authors"] != "N/A":
         authors = metadata['authors']
-        if len(authors) > 100:  # Truncate long author lists
+        if len(authors) > 100:
             authors = authors[:100] + "..."
-        meta_lines.append(f"Authors: {authors}")
+        meta_lines.append(f"👥 {authors}")
     if metadata.get("year"):
-        meta_lines.append(f"Year: {metadata['year']}")
+        meta_lines.append(f"📅 {metadata['year']}")
     if metadata.get("primary_category") and metadata["primary_category"] != "N/A":
-        meta_lines.append(f"Category: {metadata['primary_category']}")
-    return " | ".join(meta_lines)
+        meta_lines.append(f"🏷️ {metadata['primary_category']}")
+    return " • ".join(meta_lines)
 
 def format_docs(docs):
     """Format documents with clear structure and metadata."""
     if not docs:
-        return "No relevant documents found in the database. The system may need more data or a rebuilt index."
+        return "No relevant documents found in the database."
     
     formatted_chunks = []
     for idx, doc in enumerate(docs, start=1):
         meta_str = _format_metadata(doc.metadata)
         content = doc.page_content.strip()
         
-        # Truncate very long content
         if len(content) > 1000:
             content = content[:1000] + "..."
         
-        formatted_chunk = f"[Document {idx}]\n{meta_str}\n\nContent: {content}"
+        formatted_chunk = f"[Document {idx}]\n{meta_str}\n\n{content}"
         formatted_chunks.append(formatted_chunk)
     
     return "\n\n" + "="*80 + "\n\n".join(formatted_chunks)
@@ -240,18 +620,12 @@ def build_chain():
         use_rerank=use_rerank,
     )
     
-    # Improved retrieval with error handling
     def retrieval_with_logging(q):
         try:
             docs = retriever.get_relevant_documents(q)
-            formatted = format_docs(docs)
-            # Debug info in sidebar
-            with st.sidebar:
-                st.write(f"📊 Retrieved {len(docs)} documents")
-            return formatted
+            return format_docs(docs)
         except Exception as e:
-            st.error(f"Retrieval error: {e}")
-            return "Error retrieving documents. Please try again or rebuild the index."
+            return f"Error retrieving documents: {e}"
     
     retrieval_runnable = RunnableLambda(retrieval_with_logging)
     chain = {"context": retrieval_runnable, "question": RunnablePassthrough()} | prompt | llm
@@ -260,58 +634,276 @@ def build_chain():
 # Initialize session state
 if "messages" not in st.session_state:
     st.session_state["messages"] = []
+    st.session_state["show_welcome"] = True
 
-# Display index stats
-try:
-    index_stats = vectorstore.index.ntotal if hasattr(vectorstore, 'index') else "Unknown"
-    st.sidebar.info(f"📚 Vector store contains {index_stats} embeddings")
-except:
-    pass
+# Welcome message with suggestions
+if st.session_state.get("show_welcome", False):
+    st.markdown("""
+    <div class="welcome-message">
+        <h2>👋 Welcome to Research Assistant!</h2>
+        <p>I'm your AI-powered research companion. Ask me anything about Machine Learning papers!</p>
+        <div style="margin-top: 20px;">
+            <span class="suggestion-chip">🔍 Find papers on transformers</span>
+            <span class="suggestion-chip">💡 Explain attention mechanism</span>
+            <span class="suggestion-chip">📊 Compare CNN vs RNN</span>
+            <span class="suggestion-chip">🎯 Latest in reinforcement learning</span>
+        </div>
+    </div>
+    """, unsafe_allow_html=True)
+    st.session_state["show_welcome"] = False
 
-query = st.chat_input("Ask me something...")
+# Helper functions
+def is_casual_conversation(query_text):
+    """Check if the query is a greeting or casual conversation."""
+    query_lower = query_text.lower().strip()
+    greetings = ["hi", "hello", "hey", "good morning", "good afternoon", "good evening", 
+                 "hola", "greetings", "howdy", "yo", "sup", "what's up", "whats up"]
+    casual_patterns = [
+        "how are you", "how r u", "how do you do", "what's up", "whats up",
+        "who are you", "what are you", "what is your name", "your name",
+        "what can you do", "help me", "can you help", "thank you", "thanks",
+        "bye", "goodbye", "see you", "nice to meet you", "pleasure"
+    ]
+    
+    if query_lower in greetings:
+        return True
+    for pattern in casual_patterns:
+        if pattern in query_lower:
+            return True
+    return False
 
-if query:
-    with st.spinner("🔍 Searching papers and generating response..."):
-        rag_chain, adv_retriever = build_chain()
+def get_casual_response(query_text):
+    """Generate appropriate response for casual conversation."""
+    query_lower = query_text.lower().strip()
+    
+    if any(word in query_lower for word in ["hi", "hello", "hey", "hola", "howdy", "yo"]):
+        return "Hello! 👋 I'm your AI Research Assistant for Machine Learning papers. How can I help you today?"
+    if "good morning" in query_lower:
+        return "Good morning! ☀️ Ready to explore some ML research? What interests you today?"
+    if "good afternoon" in query_lower:
+        return "Good afternoon! 🌤️ Let's dive into some research! What would you like to learn about?"
+    if "good evening" in query_lower:
+        return "Good evening! 🌙 I'm here to help with ML research. What topic interests you?"
+    if any(phrase in query_lower for phrase in ["how are you", "how r u", "how do you do"]):
+        return "I'm doing great, thanks! 😊 Ready to help you explore ML research. What's on your mind?"
+    if any(phrase in query_lower for phrase in ["who are you", "what are you", "your name"]):
+        return "I'm an AI Research Assistant specialized in Machine Learning! 🤖 I help you find papers, explain concepts, and answer research questions. What would you like to know?"
+    if any(phrase in query_lower for phrase in ["what can you do", "help me", "can you help"]):
+        return """I can help you with:
         
-        # Get documents first for display
-        try:
-            docs = adv_retriever.get_relevant_documents(query)
-            
-            # Show retrieval success/failure
-            if not docs:
-                st.warning("⚠️ No documents found matching your query. Try broader terms or rebuild the index.")
+🔍 **Finding research papers** on specific ML topics  
+📚 **Explaining ML concepts** from published research  
+💡 **Answering questions** about techniques and methods  
+🎓 **Exploring** the latest ML research developments
+
+Try asking:
+- "Find papers on deep learning"
+- "What is transfer learning?"
+- "Explain adversarial training"
+
+What interests you?"""
+    if any(word in query_lower for word in ["thank you", "thanks", "thx"]):
+        return "You're welcome! 😊 Happy to help! Let me know if you have other questions."
+    if any(word in query_lower for word in ["bye", "goodbye", "see you"]):
+        return "Goodbye! 👋 Come back anytime for ML research help. Happy learning!"
+    
+    return "I'm here to help with Machine Learning research! 😊 Ask me about any ML topics or papers."
+
+# Chat input
+query = st.chat_input("💬 Ask me anything about ML research...")
+
+# Display chat history
+for i, msg in enumerate(st.session_state["messages"]):
+    # Show user message
+    st.chat_message("user", avatar="👤").write(msg["query"])
+    
+    # Show assistant response if available
+    if msg.get("answer") is not None:
+        with st.chat_message("assistant", avatar="🤖"):
+            st.write(msg["answer"])
+            if msg.get("context") and len(msg["context"]) > 0:
+                with st.expander(f"📄 View {len(msg['context'])} Retrieved Documents", expanded=False):
+                    for idx, doc in enumerate(msg["context"], 1):
+                        st.markdown(f"**📎 Document {idx}**")
+                        st.caption(_format_metadata(doc.metadata))
+                        st.text_area(
+                            f"Content {idx}", 
+                            doc.page_content[:800] + ("..." if len(doc.page_content) > 800 else ""),
+                            height=150,
+                            key=f"doc_{i}_{idx}",
+                            disabled=True
+                        )
+                        if idx < len(msg["context"]):
+                            st.markdown("---")
+    else:
+        # Answer is being generated - show thinking indicator
+        with st.chat_message("assistant", avatar="🤖"):
+            thinking_placeholder = st.empty()
+            thinking_placeholder.markdown('<p class="thinking">🔍 Searching research papers...</p>', unsafe_allow_html=True)
             
-            # Generate answer
-            answer = rag_chain.invoke(query)
-            answer_text = answer.content if hasattr(answer, "content") else str(answer)
+            # Check if casual conversation
+            if is_casual_conversation(msg["query"]):
+                casual_response = get_casual_response(msg["query"])
+                
+                # Smooth streaming effect
+                response_placeholder = st.empty()
+                full_response = ""
+                words = casual_response.split()
+                
+                for word in words:
+                    full_response += word + " "
+                    response_placeholder.markdown(full_response)
+                    time.sleep(0.02)
+                
+                st.session_state["messages"][i]["answer"] = casual_response
+                st.rerun()
             
-        except Exception as e:
-            msg = str(e)
-            if "models/" in msg and "not found" in msg.lower():
-                st.error("Selected Gemini model not found or unsupported. Try 'gemini-1.5-pro-latest' or check your model name in the sidebar.")
             else:
-                st.error(f"LLM error: {e}")
-            st.stop()
+                # Research question - full RAG pipeline
+                rag_chain, adv_retriever = build_chain()
+                
+                docs = []
+                answer_text = ""
+                error_occurred = False
+            
+                try:
+                    docs = adv_retriever.get_relevant_documents(msg["query"])
+                    
+                    if not docs:
+                        answer_text = """I couldn't find any relevant research papers in the database that match your query.
+
+**💡 Suggestions:**
+- Try using broader or different search terms
+- Check the spelling of technical terms
+- The database may not contain papers on this specific topic
+- Consider rebuilding the index with more data
+
+The current database focuses on ArXiv ML papers, but may not cover all research areas comprehensively."""
+                    else:
+                        thinking_placeholder.markdown('<p class="thinking">🧠 Analyzing documents...</p>', unsafe_allow_html=True)
+                        
+                        # Check relevance
+                        formatted_context = format_docs(docs)
+                        relevance_check_chain = {"context": RunnablePassthrough(), "question": RunnablePassthrough()} | relevance_prompt | llm
+                        relevance_result = relevance_check_chain.invoke({"context": formatted_context, "question": msg["query"]})
+                        relevance_text = relevance_result.content if hasattr(relevance_result, "content") else str(relevance_result)
+                        
+                        if "IRRELEVANT" in relevance_text.strip().upper():
+                            answer_text = f"""I found {len(docs)} documents in the database, but they don't contain relevant information about your question.
+
+**📋 Retrieved topics:**
+- {docs[0].metadata.get('title', 'Various topics') if docs else 'N/A'}
+
+**💡 Suggestions:**
+- Try rephrasing with different keywords
+- Use more specific technical terms
+- Search for related concepts or broader topics
+- The database may not have papers specifically on this topic
+
+I can only provide answers based on the ArXiv papers in the database."""
+                        else:
+                            # Generate answer with streaming
+                            thinking_placeholder.markdown('<p class="thinking">✍️ Generating response...</p>', unsafe_allow_html=True)
+                            answer = rag_chain.invoke(msg["query"])
+                            answer_text = answer.content if hasattr(answer, "content") else str(answer)
+                    
+                except Exception as e:
+                    error_occurred = True
+                    msg_err = str(e)
+                    if "models/" in msg_err and "not found" in msg_err.lower():
+                        answer_text = "⚠️ Selected model not found. Try a different model in the sidebar."
+                    else:
+                        answer_text = f"⚠️ An error occurred: {e}\n\nPlease try again or rebuild the index."
+                
+                # Clear thinking and display response with streaming
+                thinking_placeholder.empty()
+                
+                # Stream response
+                import re
+                response_placeholder = st.empty()
+                parts = re.split(r'(\n\n|(?<=[.!?])\s+)', answer_text)
+                
+                full_response = ""
+                for part in parts:
+                    full_response += part
+                    response_placeholder.markdown(full_response)
+                    time.sleep(0.03)
+                
+                # Update session state
+                st.session_state["messages"][i]["answer"] = answer_text
+                st.session_state["messages"][i]["context"] = docs
+                
+                # Show retrieved documents
+                if docs:
+                    with st.expander(f"📄 View {len(docs)} Retrieved Documents", expanded=False):
+                        for idx, doc in enumerate(docs, 1):
+                            st.markdown(f"**📎 Document {idx}**")
+                            st.caption(_format_metadata(doc.metadata))
+                            st.text_area(
+                                f"Content {idx}", 
+                                doc.page_content[:800] + ("..." if len(doc.page_content) > 800 else ""),
+                                height=150,
+                                key=f"new_doc_{i}_{idx}",
+                                disabled=True
+                            )
+                            if idx < len(docs):
+                                st.markdown("---")
+                
+                st.rerun()
+
+# Process new query
+if query:
+    # Add message to session state immediately
+    st.session_state["messages"].append({
+        "query": query,
+        "answer": None,
+        "context": []
+    })
+    
+    # Force rerun to show the user message immediately
+    st.rerun()
+
+# Footer with tips - only show if there are messages
+if len(st.session_state["messages"]) > 0:
+    st.markdown("---")
+    with st.expander("💡 Tips for Better Results", expanded=False):
+        col1, col2 = st.columns(2)
         
-        # Save to session state
-        st.session_state["messages"].append({
-            "query": query,
-            "answer": answer_text,
-            "context": docs
-        })
-
-# Display conversation history
-for msg in st.session_state["messages"]:
-    st.chat_message("user").write(msg["query"])
-    with st.chat_message("assistant"):
-        st.write(msg["answer"])
-        with st.expander(f"📄 View {len(msg['context'])} Retrieved Documents"):
-            if not msg["context"]:
-                st.info("No documents were retrieved for this query.")
-            for i, doc in enumerate(msg["context"], 1):
-                st.markdown(f"### Document {i}")
-                st.write(doc.page_content[:500] + ("..." if len(doc.page_content) > 500 else ""))
-                if doc.metadata:
-                    st.caption(_format_metadata(doc.metadata))
-                st.divider()
+        with col1:
+            st.markdown("""
+            **🎯 Asking Better Questions**
+            
+            ✅ Use specific ML terminology  
+            ✅ Mention techniques or methods  
+            ✅ Ask for comparisons  
+            ✅ Reference specific problems  
+            
+            **Examples:**
+            - "Papers on transformer architecture"
+            - "Compare CNNs vs Vision Transformers"
+            - "Explain BERT training methodology"
+            """)
+        
+        with col2:
+            st.markdown("""
+            **📚 Understanding Responses**
+            
+            ✅ All answers from actual papers  
+            ✅ View source documents anytime  
+            ✅ Check relevance of results  
+            ✅ Adjust settings if needed  
+            
+            **⚡ Advanced Tips:**
+            - Use sidebar filters (year, category)
+            - Adjust retrieval settings
+            - Try different LLM providers
+            - Rebuild index for fresh data
+            """)
+
+# Add a "Clear Chat" button at the bottom of sidebar
+with st.sidebar:
+    st.markdown("---")
+    if st.button("🗑️ Clear Chat History", use_container_width=True):
+        st.session_state["messages"] = []
+        st.session_state["show_welcome"] = True
+        st.rerun()