Update app.py

app.py

@@ -1,23 +1,29 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from PyPDF2 import PdfReader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
+from langchain_community.vectorstores import FAISS
+from langchain_core.prompts import PromptTemplate
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 import time
-from datetime import datetime, timezone, timedelta
-import html as html_lib
-
-# ─── TIMEZONE (IST = UTC+5:30) ────────────────────────────────────────────────
-IST = timezone(timedelta(hours=5, minutes=30))
-
-def get_ist_time():
-    return datetime.now(IST).strftime("%H:%M")
+import base64
 
 st.set_page_config(
-    page_title="TurboChat AI",
-    page_icon="⚡",
-    layout="centered",
+    page_title="QueryDocs AI",
+    page_icon="📚",
+    layout="wide",
     initial_sidebar_state="expanded"
 )
 
+# ─── HELPERS ──────────────────────────────────────────────────────────────────
+def img_to_base64(path):
+    try:
+        with open(path, "rb") as f:
+            return base64.b64encode(f.read()).decode()
+    except:
+        return None
+
 # ─── GLOBAL CSS ───────────────────────────────────────────────────────────────
 st.markdown("""
 <style>
@@ -37,241 +43,292 @@ html, body, [data-testid="stAppViewContainer"] {
 footer, #MainMenu { visibility: hidden; }
 
 /* ── Header ── */
-.chat-header {
-    text-align: center;
-    padding: 20px 0 10px;
+.app-header {
+    display: flex; align-items: center; gap: 14px;
+    padding: 16px 0 14px;
     border-bottom: 1px solid #21262d;
     margin-bottom: 20px;
+    animation: fadeInDown 0.5s ease both;
 }
-.chat-title {
+.app-title {
     font-family: 'Orbitron', monospace;
-    font-size: 1.8rem;
-    font-weight: 900;
-    color: #fff;
-    letter-spacing: 2px;
+    font-size: 1.6rem; font-weight: 900;
+    color: #fff; letter-spacing: 2px;
 }
-.chat-title span { color: #58a6ff; }
-.chat-sub {
+.app-title span { color: #58a6ff; }
+.app-sub {
     font-family: 'Share Tech Mono', monospace;
-    font-size: 0.7rem;
-    color: #8b949e;
-    letter-spacing: 3px;
-    margin-top: 4px;
+    font-size: 0.68rem; color: #8b949e;
+    letter-spacing: 3px; margin-top: 3px;
 }
 
-/* ── Chat bubbles ── */
-.chat-wrap {
-    display: flex;
-    flex-direction: column;
-    gap: 14px;
-    padding: 10px 0;
-}
-.msg-user {
-    display: flex;
-    justify-content: flex-end;
-    animation: slideInRight 0.3s ease both;
-}
-.msg-ai {
-    display: flex;
-    justify-content: flex-start;
-    animation: slideInLeft 0.3s ease both;
+/* ── Profile card in sidebar ── */
+.profile-card {
+    text-align: center;
+    padding: 20px 0 16px;
+    border-bottom: 1px solid #21262d;
+    margin-bottom: 16px;
+    animation: fadeInDown 0.5s ease both;
+}
+.profile-avatar {
+    width: 72px; height: 72px;
+    border-radius: 50%;
+    overflow: hidden;
+    border: 2px solid #58a6ff;
+    box-shadow: 0 0 0 3px rgba(88,166,255,0.15),
+                0 0 20px rgba(88,166,255,0.2);
+    margin: 0 auto 10px;
+    animation: pulse-av 2.5s ease-in-out infinite;
+}
+.profile-avatar img {
+    width: 100%; height: 100%;
+    object-fit: cover; border-radius: 50%;
+}
+@keyframes pulse-av {
+    0%,100%{box-shadow:0 0 0 3px rgba(88,166,255,0.15),0 0 20px rgba(88,166,255,0.2);}
+    50%    {box-shadow:0 0 0 4px rgba(88,166,255,0.3),0 0 30px rgba(88,166,255,0.35);}
+}
+.profile-name {
+    font-family: 'Orbitron', monospace;
+    font-size: 0.82rem; font-weight: 700;
+    color: #fff; letter-spacing: 1px;
+}
+.profile-role {
+    font-family: 'Share Tech Mono', monospace;
+    font-size: 0.64rem; color: #58a6ff;
+    letter-spacing: 2px; margin-top: 3px;
+}
+.profile-links {
+    display: flex; justify-content: center;
+    gap: 8px; margin-top: 10px; flex-wrap: wrap;
+}
+.p-link {
+    font-family: 'Share Tech Mono', monospace;
+    font-size: 0.62rem; color: #8b949e !important;
+    text-decoration: none !important;
+    background: #0d1117;
+    border: 1px solid #30363d;
+    padding: 3px 10px; border-radius: 20px;
+    transition: all 0.25s;
+}
+.p-link:hover { color: #58a6ff !important; border-color: #58a6ff; }
+
+/* ── Upload zone ── */
+.upload-card {
+    background: #161b22;
+    border: 2px dashed #30363d;
+    border-radius: 12px;
+    padding: 32px;
+    text-align: center;
+    transition: all 0.3s ease;
+    animation: fadeInUp 0.5s ease both;
+}
+.upload-card:hover { border-color: #58a6ff; }
+.upload-icon { font-size: 2.5rem; margin-bottom: 10px; }
+.upload-title {
+    font-family: 'Orbitron', monospace;
+    font-size: 0.95rem; color: #fff; margin-bottom: 6px;
 }
+.upload-sub { font-size: 0.85rem; color: #8b949e; }
+
+/* ── PDF info banner ── */
+.pdf-banner {
+    background: #161b22;
+    border: 1px solid #21262d;
+    border-left: 3px solid #58a6ff;
+    border-radius: 8px;
+    padding: 12px 18px;
+    display: flex; align-items: center; gap: 12px;
+    margin-bottom: 16px;
+    animation: fadeInUp 0.4s ease both;
+}
+.pdf-name {
+    font-weight: 700; font-size: 0.95rem; color: #fff;
+}
+.pdf-meta {
+    font-family: 'Share Tech Mono', monospace;
+    font-size: 0.68rem; color: #8b949e; margin-top: 2px;
+}
+
+/* ── Chat bubbles ── */
+.chat-wrap { display: flex; flex-direction: column; gap: 16px; padding: 8px 0; }
+.msg-user  { display: flex; justify-content: flex-end;  animation: slideInR 0.3s ease both; }
+.msg-ai    { display: flex; justify-content: flex-start; animation: slideInL 0.3s ease both; }
 .bubble-user {
     background: linear-gradient(135deg, #1f6feb, #388bfd);
-    color: #fff;
-    padding: 12px 18px;
+    color: #fff; padding: 12px 18px;
     border-radius: 18px 18px 4px 18px;
-    max-width: 78%;
-    font-size: 0.97rem;
-    line-height: 1.6;
+    max-width: 75%; font-size: 0.97rem; line-height: 1.6;
     box-shadow: 0 4px 15px rgba(31,111,235,0.25);
 }
 .bubble-ai {
-    background: #161b22;
-    color: #c9d1d9;
-    padding: 12px 18px;
+    background: #161b22; color: #c9d1d9;
+    padding: 14px 18px;
     border-radius: 18px 18px 18px 4px;
-    max-width: 78%;
-    font-size: 0.97rem;
-    line-height: 1.7;
+    max-width: 80%; font-size: 0.95rem; line-height: 1.75;
     border: 1px solid #30363d;
     box-shadow: 0 4px 15px rgba(0,0,0,0.3);
-    position: relative;
 }
-.bubble-ai::before {
-    content: '⚡';
-    position: absolute;
-    top: -10px; left: 12px;
-    font-size: 0.7rem;
-    background: #0d1117;
-    padding: 0 4px;
-    border-radius: 50%;
+.bubble-ai .answer-label {
+    font-family: 'Share Tech Mono', monospace;
+    font-size: 0.65rem; color: #58a6ff;
+    letter-spacing: 2px; margin-bottom: 8px;
+    text-transform: uppercase;
 }
 .msg-meta {
     font-family: 'Share Tech Mono', monospace;
-    font-size: 0.62rem;
-    color: #8b949e;
-    margin-top: 4px;
-    padding: 0 6px;
+    font-size: 0.62rem; color: #484f58;
+    margin-top: 4px; padding: 0 6px;
 }
 
-/* ── Typing animation ── */
-.typing-indicator {
-    display: flex;
-    align-items: center;
-    gap: 5px;
-    padding: 14px 18px;
-    background: #161b22;
-    border: 1px solid #30363d;
-    border-radius: 18px 18px 18px 4px;
-    width: fit-content;
-    animation: slideInLeft 0.3s ease both;
+.answer-text {
+    white-space: pre-wrap;
+    word-break: break-word;
+    line-height: 1.75;
+    font-size: 0.95rem;
 }
-.typing-dot {
-    width: 7px; height: 7px;
-    background: #58a6ff;
-    border-radius: 50%;
-    animation: bounce 1.2s ease-in-out infinite;
+
+/* ── Source chunks ── */
+.sources-wrap {
+    margin-top: 10px;
+    border-top: 1px solid #21262d;
+    padding-top: 10px;
 }
-.typing-dot:nth-child(2) { animation-delay: 0.2s; }
-.typing-dot:nth-child(3) { animation-delay: 0.4s; }
-@keyframes bounce {
-    0%,60%,100% { transform: translateY(0); opacity:0.4; }
-    30%          { transform: translateY(-8px); opacity:1; }
+.source-label {
+    font-family: 'Share Tech Mono', monospace;
+    font-size: 0.62rem; color: #8b949e;
+    letter-spacing: 2px; margin-bottom: 6px;
 }
-
-/* ── Stats bar ── */
-.stats-bar {
-    display: flex;
-    gap: 16px;
-    flex-wrap: wrap;
-    padding: 8px 12px;
-    background: #161b22;
-    border: 1px solid #21262d;
-    border-radius: 8px;
-    margin-bottom: 12px;
+.source-chip {
+    display: inline-block;
+    background: rgba(88,166,255,0.08);
+    border: 1px solid rgba(88,166,255,0.2);
+    color: #58a6ff; padding: 3px 10px;
+    border-radius: 4px; font-size: 0.72rem;
     font-family: 'Share Tech Mono', monospace;
-    font-size: 0.68rem;
-    color: #8b949e;
+    margin: 3px 3px; cursor: pointer;
 }
-.stat-item { display: flex; align-items: center; gap: 5px; }
-.stat-val  { color: #58a6ff; font-weight: 600; }
 
-/* ── Input area ── */
-div[data-testid="column"]:first-child .stTextInput > div > div > input {
+/* ── Typing indicator ── */
+.typing-wrap { display: flex; justify-content: flex-start; }
+.typing-box {
+    display: flex; align-items: center; gap: 5px;
+    padding: 14px 18px; background: #161b22;
+    border: 1px solid #30363d;
+    border-radius: 18px 18px 18px 4px;
+}
+.t-dot {
+    width: 7px; height: 7px; background: #58a6ff;
+    border-radius: 50%;
+    animation: tdot 1.2s ease-in-out infinite;
+}
+.t-dot:nth-child(2){animation-delay:0.2s;}
+.t-dot:nth-child(3){animation-delay:0.4s;}
+@keyframes tdot{0%,60%,100%{transform:translateY(0);opacity:0.4;}30%{transform:translateY(-8px);opacity:1;}}
+
+/* ── Input ── */
+.stTextInput > div > div > input {
     background: #161b22 !important;
     border: 1px solid #30363d !important;
-    border-right: none !important;
-    border-radius: 12px 0 0 12px !important;
+    border-radius: 8px !important;
     color: #c9d1d9 !important;
     font-family: 'Rajdhani', sans-serif !important;
     font-size: 0.97rem !important;
-    padding: 12px 18px !important;
-    height: 48px !important;
 }
-div[data-testid="column"]:first-child .stTextInput > div > div > input:focus {
+.stTextInput > div > div > input:focus {
     border-color: #58a6ff !important;
-    box-shadow: none !important;
-    outline: none !important;
-}
-/* Send button column */
-div[data-testid="column"]:last-child .stButton > button {
-    background: linear-gradient(135deg, #1f6feb, #388bfd) !important;
-    color: #fff !important;
-    border: 1px solid #1f6feb !important;
-    border-left: none !important;
-    border-radius: 0 12px 12px 0 !important;
-    height: 48px !important;
-    width: 100% !important;
-    font-size: 1.1rem !important;
-    letter-spacing: 0 !important;
-    box-shadow: 0 0 14px rgba(31,111,235,0.3) !important;
-    transition: all 0.2s ease !important;
-    margin-top: 0 !important;
-    padding: 0 !important;
-}
-div[data-testid="column"]:last-child .stButton > button:hover {
-    box-shadow: 0 0 24px rgba(31,111,235,0.55) !important;
-    color: #fff !important;
-}
-/* Remove gap between columns so they touch */
-div[data-testid="column"]:first-child { padding-right: 0 !important; }
-div[data-testid="column"]:last-child  { padding-left: 0 !important; }
+    box-shadow: 0 0 12px rgba(88,166,255,0.15) !important;
+}
 
 /* ── Buttons ── */
 .stButton > button {
-    background: #21262d !important;
-    color: #c9d1d9 !important;
-    border: 1px solid #30363d !important;
-    border-radius: 8px !important;
+    background: #21262d !important; color: #c9d1d9 !important;
+    border: 1px solid #30363d !important; border-radius: 8px !important;
     font-family: 'Share Tech Mono', monospace !important;
-    font-size: 0.75rem !important;
-    letter-spacing: 1px !important;
-    transition: all 0.25s ease !important;
+    font-size: 0.73rem !important; letter-spacing: 1px !important;
+    transition: all 0.25s !important;
 }
 .stButton > button:hover {
     background: #30363d !important;
-    border-color: #58a6ff !important;
-    color: #58a6ff !important;
+    border-color: #58a6ff !important; color: #58a6ff !important;
+}
+.send-btn > button {
+    background: linear-gradient(135deg,#1f6feb,#388bfd) !important;
+    color: #fff !important; border: none !important;
+    box-shadow: 0 0 16px rgba(31,111,235,0.3) !important;
+}
+.send-btn > button:hover {
+    box-shadow: 0 0 28px rgba(31,111,235,0.55) !important;
+    transform: translateY(-2px) !important; color: #fff !important;
 }
 
-
-
-/* ── Sidebar inputs ── */
-[data-testid="stSidebar"] textarea,
-[data-testid="stSidebar"] input {
-    background: #0d1117 !important;
-    border: 1px solid #30363d !important;
-    color: #c9d1d9 !important;
-    border-radius: 6px !important;
+/* ── Stats ── */
+.stat-row {
+    display: flex; gap: 10px; flex-wrap: wrap;
+    margin-bottom: 14px;
 }
-[data-testid="stSidebar"] .stSelectbox > div > div {
-    background: #0d1117 !important;
-    border: 1px solid #30363d !important;
-    color: #c9d1d9 !important;
+.stat-box {
+    flex: 1; min-width: 80px;
+    background: #161b22; border: 1px solid #21262d;
+    border-top: 2px solid #58a6ff; border-radius: 6px;
+    padding: 10px 12px; text-align: center;
+}
+.stat-num {
+    font-family: 'Orbitron', monospace;
+    font-size: 1.2rem; font-weight: 900; color: #58a6ff;
+}
+.stat-lbl {
+    font-family: 'Share Tech Mono', monospace;
+    font-size: 0.6rem; color: #8b949e;
+    letter-spacing: 1px; margin-top: 3px;
 }
 
-/* ── Welcome card ── */
+/* ── Welcome ── */
 .welcome-card {
-    background: #161b22;
-    border: 1px solid #21262d;
-    border-radius: 12px;
-    padding: 28px;
-    text-align: center;
-    margin: 20px 0;
+    background: #161b22; border: 1px solid #21262d;
+    border-radius: 12px; padding: 30px;
+    text-align: center; margin: 10px 0;
     animation: fadeInUp 0.6s ease both;
 }
-.welcome-icon { font-size: 2.5rem; margin-bottom: 10px; }
-.welcome-title {
+.wc-icon { font-size: 2.8rem; margin-bottom: 12px; }
+.wc-title {
     font-family: 'Orbitron', monospace;
     font-size: 1rem; color: #fff; margin-bottom: 8px;
 }
-.welcome-sub { font-size: 0.9rem; color: #8b949e; line-height: 1.6; }
+.wc-sub { font-size: 0.9rem; color: #8b949e; line-height: 1.65; }
 .tip-chip {
     display: inline-block;
-    background: rgba(88,166,255,0.1);
-    border: 1px solid rgba(88,166,255,0.3);
-    color: #58a6ff;
-    padding: 4px 12px; border-radius: 20px;
+    background: rgba(88,166,255,0.08);
+    border: 1px solid rgba(88,166,255,0.25);
+    color: #58a6ff; padding: 4px 12px;
+    border-radius: 20px; margin: 4px 3px;
     font-family: 'Share Tech Mono', monospace;
-    font-size: 0.68rem; margin: 4px 3px;
+    font-size: 0.67rem;
 }
 
-/* ── Animations ── */
-@keyframes slideInRight {
-    from { opacity:0; transform: translateX(20px); }
-    to   { opacity:1; transform: translateX(0); }
+/* ── File uploader ── */
+[data-testid="stFileUploader"] {
+    background: #161b22 !important;
+    border: 2px dashed #30363d !important;
+    border-radius: 10px !important;
 }
-@keyframes slideInLeft {
-    from { opacity:0; transform: translateX(-20px); }
-    to   { opacity:1; transform: translateX(0); }
+[data-testid="stFileUploader"]:hover {
+    border-color: #58a6ff !important;
 }
-@keyframes fadeInUp {
-    from { opacity:0; transform: translateY(16px); }
-    to   { opacity:1; transform: translateY(0); }
+
+/* ── Glow divider ── */
+.glow-div {
+    border: none; height: 1px;
+    background: linear-gradient(90deg,transparent,#58a6ff,transparent);
+    margin: 16px 0; box-shadow: 0 0 8px rgba(88,166,255,0.3);
 }
 
-/* scrollbar */
+/* ── Animations ── */
+@keyframes fadeInDown { from{opacity:0;transform:translateY(-16px)} to{opacity:1;transform:translateY(0)} }
+@keyframes fadeInUp   { from{opacity:0;transform:translateY(16px)}  to{opacity:1;transform:translateY(0)} }
+@keyframes slideInR   { from{opacity:0;transform:translateX(20px)}  to{opacity:1;transform:translateX(0)} }
+@keyframes slideInL   { from{opacity:0;transform:translateX(-20px)} to{opacity:1;transform:translateX(0)} }
+
 ::-webkit-scrollbar { width: 5px; }
 ::-webkit-scrollbar-track { background: #0d1117; }
 ::-webkit-scrollbar-thumb { background: #30363d; border-radius: 3px; }
@@ -279,225 +336,331 @@ div[data-testid="column"]:last-child  { padding-left: 0 !important; }
 </style>
 """, unsafe_allow_html=True)
 
-# ─── MODEL LOAD ───────────────────────────────────────────────────────────────
-@st.cache_resource
-def load_model():
-    model_name = "microsoft/Phi-3-mini-4k-instruct"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+# ─── SESSION STATE ────────────────────────────────────────────────────────────
+if "messages"    not in st.session_state: st.session_state.messages    = []
+if "vectorstore" not in st.session_state: st.session_state.vectorstore = None
+if "pdf_name"    not in st.session_state: st.session_state.pdf_name    = None
+if "pdf_pages"   not in st.session_state: st.session_state.pdf_pages   = 0
+if "pdf_chunks"  not in st.session_state: st.session_state.pdf_chunks  = 0
+if "q_count"     not in st.session_state: st.session_state.q_count     = 0
+if "last_q"      not in st.session_state: st.session_state.last_q      = ""
+
+# ─── MODEL LOADERS ────────────────────────────────────────────────────────────
+@st.cache_resource(show_spinner=False)
+def load_embeddings():
+    return HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
-    model.eval()
-    if torch.cuda.is_available():
-        model = model.to("cuda")
-    return tokenizer, model
 
-# ─── SESSION STATE ────────────────────────────────────────────────────────────
-if "messages"     not in st.session_state: st.session_state.messages     = []
-if "total_tokens" not in st.session_state: st.session_state.total_tokens = 0
-if "total_time"   not in st.session_state: st.session_state.total_time   = 0.0
-if "msg_count"    not in st.session_state: st.session_state.msg_count    = 0
-if "last_q"       not in st.session_state: st.session_state.last_q       = ""
-if "input_key"    not in st.session_state: st.session_state.input_key    = 0
+@st.cache_resource(show_spinner=False)
+def load_llm():
+    model_id  = "TinyLlama/TinyLlama-1.1B-chat-v1.0"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model     = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        low_cpu_mem_usage=True,
+        device_map="cuda" if torch.cuda.is_available() else None
+    )
+    if not torch.cuda.is_available():
+        model = model.to("cpu")
+    pipe = pipeline(
+        "text-generation", model=model, tokenizer=tokenizer,
+        max_new_tokens=512, temperature=0.3, do_sample=True,
+        pad_token_id=tokenizer.eos_token_id, repetition_penalty=1.1
+    )
+    return HuggingFacePipeline(pipeline=pipe)
+
+# ─── PDF PROCESSOR ────────────────────────────────────────────────────────────
+def process_pdf(uploaded_file):
+    reader   = PdfReader(uploaded_file)
+    raw_text = ""
+    for page in reader.pages:
+        text = page.extract_text()
+        if text:
+            raw_text += text
+
+    if not raw_text.strip():
+        raise ValueError("No readable text found. PDF may be scanned/image-based.")
+
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks   = splitter.split_text(raw_text)
+    embeddings  = load_embeddings()
+    vectorstore = FAISS.from_texts(chunks, embeddings)
+    return vectorstore, len(reader.pages), len(chunks)
+
+# ─── ANSWER FUNCTION ──────────────────────────────────────────────────────────
+def get_answer(question, vectorstore):
+    retriever     = vectorstore.as_retriever(search_kwargs={"k": 4})
+    relevant_docs = retriever.invoke(question)
+    context       = "\n\n".join([f"---\n{doc.page_content}" for doc in relevant_docs])
+    sources       = [doc.page_content[:120] + "..." for doc in relevant_docs]
+
+    prompt_template = PromptTemplate(
+        input_variables=["context", "question"],
+        template="""<|system|>
+You are QueryDocs AI, an intelligent document assistant. Use ONLY the context provided to answer the question clearly and accurately. If the answer is not in the context, say so honestly.
+<|user|>
+CONTEXT:
+{context}
+
+QUESTION:
+{question}
+<|assistant|>
+"""
+    )
+    llm    = load_llm()
+    chain  = prompt_template | llm
+    result = chain.invoke({"context": context, "question": question})
+
+    # extract only the assistant reply
+    if "<|assistant|>" in result:
+        answer = result.split("<|assistant|>")[-1].strip()
+    else:
+        answer = result.strip()
+
+    return answer, sources
 
 # ─── SIDEBAR ──────────────────────────────────────────────────────────────────
 with st.sidebar:
-    st.markdown("""
-    <div style="text-align:center;padding:16px 0 20px;">
-        <div style="font-family:'Orbitron',monospace;font-size:1rem;color:#fff;letter-spacing:2px;">⚡ TURBO<span style="color:#58a6ff;">CHAT</span></div>
-        <div style="font-family:'Share Tech Mono',monospace;font-size:0.65rem;color:#8b949e;letter-spacing:3px;margin-top:4px;">PHI-3 MINI</div>
+    # Profile card
+    img_b64 = img_to_base64("assets/NANII.png")
+    avatar  = (f'<img src="data:image/png;base64,{img_b64}" alt="Sriram">'
+               if img_b64 else
+               '<div style="width:72px;height:72px;background:#1f6feb;border-radius:50%;margin:0 auto;"></div>')
+
+    st.markdown(f"""
+    <div class="profile-card">
+        <div class="profile-avatar">{avatar}</div>
+        <div class="profile-name">SRIRAM SAI</div>
+        <div class="profile-role">AI &amp; ML ENGINEER</div>
+        <div class="profile-links">
+            <a class="p-link" href="https://github.com/sriramsai18" target="_blank">💻 GitHub</a>
+            <a class="p-link" href="https://www.linkedin.com/in/sriram-sai-laggisetti/" target="_blank">💼 LinkedIn</a>
+        </div>
     </div>
     """, unsafe_allow_html=True)
 
-    st.markdown("**🎭 SYSTEM PROMPT**")
-    system_prompt = st.text_area(
-        "",
-        value="""You, in the role of TurboChat, are an assistant designed to deliver fast, reliable, and high-quality responses while maintaining clarity and efficiency.
-You specialize in breaking down complex information into easily understandable insights without losing accuracy or depth.
-Your main goal and objective are to provide responses that are precise, logically structured, and immediately useful to the user.
-Your task is to interpret user queries correctly, detect ambiguity, apply reasoning, and present the most relevant answer in the most efficient format possible.
-To make this work as it should, you must prioritize clarity, accuracy, and usefulness over verbosity, eliminate redundant explanations, and ensure every sentence contributes meaningful value.
-When appropriate, you organize information into structured sections, lists, or steps to improve readability and comprehension. If a question is unclear, you ask a concise clarification question before answering.
-If multiple solutions exist, you present the most practical and high-impact options first.""",
-        height=100,
-        label_visibility="collapsed",
-        key="sys_prompt"
-    )
+    # PDF upload
+    st.markdown('<div style="font-family:\'Share Tech Mono\',monospace;font-size:0.7rem;color:#8b949e;letter-spacing:2px;margin-bottom:8px;">📄 UPLOAD DOCUMENT</div>', unsafe_allow_html=True)
 
-    st.markdown("---")
-    st.markdown("**⚙️ GENERATION SETTINGS**")
+    uploaded_file = st.file_uploader(
+        "Upload PDF", type=["pdf"],
+        label_visibility="collapsed"
+    )
 
-    max_tokens = st.slider("Max Tokens", 50, 300, 120, 10)
-    temperature = st.slider("Temperature", 0.1, 1.5, 0.6, 0.1,
-                            help="Higher = more creative, Lower = more focused")
-    top_p = st.slider("Top-P", 0.1, 1.0, 0.9, 0.05,
-                      help="Nucleus sampling threshold")
+    if uploaded_file:
+        if st.session_state.pdf_name != uploaded_file.name:
+            with st.spinner("🔍 Processing PDF..."):
+                try:
+                    vs, pages, chunks = process_pdf(uploaded_file)
+                    st.session_state.vectorstore = vs
+                    st.session_state.pdf_name    = uploaded_file.name
+                    st.session_state.pdf_pages   = pages
+                    st.session_state.pdf_chunks  = chunks
+                    st.session_state.messages    = []
+                    st.session_state.q_count     = 0
+                    st.success("✅ PDF ready!")
+                except Exception as e:
+                    st.error(f"❌ {str(e)}")
 
     st.markdown("---")
-    st.markdown("**📊 SESSION STATS**")
-    st.markdown(f"""
-    <div style="font-family:'Share Tech Mono',monospace;font-size:0.72rem;color:#8b949e;line-height:2;">
-        💬 Messages &nbsp;&nbsp;<span style="color:#58a6ff;">{st.session_state.msg_count}</span><br>
-        🔤 Tokens &nbsp;&nbsp;&nbsp;&nbsp;<span style="color:#58a6ff;">{st.session_state.total_tokens}</span><br>
-        ⏱️ Total time &nbsp;<span style="color:#58a6ff;">{st.session_state.total_time:.1f}s</span>
-    </div>
-    """, unsafe_allow_html=True)
 
-    st.markdown("---")
+    # Stats
+    if st.session_state.pdf_name:
+        st.markdown(f"""
+        <div style="font-family:'Share Tech Mono',monospace;font-size:0.7rem;color:#8b949e;margin-bottom:10px;letter-spacing:2px;">📊 DOCUMENT STATS</div>
+        <div style="display:flex;flex-direction:column;gap:6px;">
+            <div style="background:#0d1117;border:1px solid #21262d;border-radius:6px;padding:8px 12px;font-family:'Share Tech Mono',monospace;font-size:0.7rem;">
+                📄 <span style="color:#c9d1d9;">{st.session_state.pdf_name[:22]}{'...' if len(st.session_state.pdf_name)>22 else ''}</span>
+            </div>
+            <div style="display:flex;gap:6px;">
+                <div style="flex:1;background:#0d1117;border:1px solid #21262d;border-top:2px solid #58a6ff;border-radius:6px;padding:8px;text-align:center;">
+                    <div style="font-family:'Orbitron',monospace;font-size:1rem;color:#58a6ff;">{st.session_state.pdf_pages}</div>
+                    <div style="font-family:'Share Tech Mono',monospace;font-size:0.58rem;color:#8b949e;">PAGES</div>
+                </div>
+                <div style="flex:1;background:#0d1117;border:1px solid #21262d;border-top:2px solid #58a6ff;border-radius:6px;padding:8px;text-align:center;">
+                    <div style="font-family:'Orbitron',monospace;font-size:1rem;color:#58a6ff;">{st.session_state.pdf_chunks}</div>
+                    <div style="font-family:'Share Tech Mono',monospace;font-size:0.58rem;color:#8b949e;">CHUNKS</div>
+                </div>
+                <div style="flex:1;background:#0d1117;border:1px solid #21262d;border-top:2px solid #58a6ff;border-radius:6px;padding:8px;text-align:center;">
+                    <div style="font-family:'Orbitron',monospace;font-size:1rem;color:#58a6ff;">{st.session_state.q_count}</div>
+                    <div style="font-family:'Share Tech Mono',monospace;font-size:0.58rem;color:#8b949e;">ASKED</div>
+                </div>
+            </div>
+        </div>
+        """, unsafe_allow_html=True)
+        st.markdown("---")
+
+    # Clear button
     if st.button("🗑️ CLEAR CHAT", use_container_width=True):
-        st.session_state.messages     = []
-        st.session_state.total_tokens = 0
-        st.session_state.total_time   = 0.0
-        st.session_state.msg_count    = 0
+        st.session_state.messages = []
+        st.session_state.q_count  = 0
         st.rerun()
 
-    st.markdown("""
-    <div style="text-align:center;margin-top:16px;font-family:'Share Tech Mono',monospace;font-size:0.62rem;color:#484f58;">
-        Built by <a href="https://github.com/sriramsai18" style="color:#58a6ff;text-decoration:none;">Sriram Sai</a>
-    </div>
-    """, unsafe_allow_html=True)
+    # New PDF button
+    if st.button("📄 LOAD NEW PDF", use_container_width=True):
+        st.session_state.vectorstore = None
+        st.session_state.pdf_name    = None
+        st.session_state.pdf_pages   = 0
+        st.session_state.pdf_chunks  = 0
+        st.session_state.messages    = []
+        st.session_state.q_count     = 0
+        st.rerun()
 
-# ─── HEADER ───────────────────────────────────────────────────────────────────
+# ─── MAIN AREA ────────────────────────────────────────────────────────────────
+# Header
 st.markdown("""
-<div class="chat-header">
-    <div class="chat-title">TURBO<span>CHAT</span> ⚡</div>
-    <div class="chat-sub">POWERED BY LLM's</div>
+<div class="app-header">
+    <div>
+        <div class="app-title">QUERY<span>DOCS</span> AI 📚</div>
+        <div class="app-sub">INTELLIGENT DOCUMENT Q&amp;A · RAG PIPELINE </div>
+    </div>
 </div>
 """, unsafe_allow_html=True)
 
-# ─── CHAT HISTORY ─────────────────────────────────────────────────────────────
-if not st.session_state.messages:
+# PDF banner (when loaded)
+if st.session_state.pdf_name:
+    st.markdown(f"""
+    <div class="pdf-banner">
+        <span style="font-size:1.4rem;">📄</span>
+        <div>
+            <div class="pdf-name">{st.session_state.pdf_name}</div>
+            <div class="pdf-meta">{st.session_state.pdf_pages} pages · {st.session_state.pdf_chunks} chunks · ready to query</div>
+        </div>
+        <span style="margin-left:auto;background:rgba(88,166,255,0.1);border:1px solid rgba(88,166,255,0.3);
+              color:#58a6ff;padding:4px 12px;border-radius:20px;
+              font-family:'Share Tech Mono',monospace;font-size:0.65rem;">● ACTIVE</span>
+    </div>
+    """, unsafe_allow_html=True)
+
+# Chat area
+if not st.session_state.vectorstore:
     st.markdown("""
     <div class="welcome-card">
-        <div class="welcome-icon">⚡</div>
-        <div class="welcome-title">TURBOCHAT IS READY</div>
-        <div class="welcome-sub">
-            Ask me anything — code, concepts, advice, or just chat.<br>
-            Powered by LLM.
+        <div class="wc-icon">📚</div>
+        <div class="wc-title">WELCOME TO QUERYDOCS AI</div>
+        <div class="wc-sub">
+            Upload any PDF document from the sidebar and start asking questions.<br>
+            Powered by RAG pipeline, context-aware answers.
         </div>
         <br>
-        <span class="tip-chip">💡 Try: Explain neural networks</span>
-        <span class="tip-chip">💡 Try: Write a Python function</span>
-        <span class="tip-chip">💡 Try: What is RAG?</span>
+        <span class="tip-chip">📋 Legal documents</span>
+        <span class="tip-chip">📊 Research papers</span>
+        <span class="tip-chip">📖 Study material</span>
+        <span class="tip-chip">📝 Reports</span>
     </div>
     """, unsafe_allow_html=True)
+
 else:
-    chat_html = '<div class="chat-wrap">'
-    for msg in st.session_state.messages:
-        safe_content = html_lib.escape(msg["content"])
-        if msg["role"] == "user":
-            chat_html += f"""
-            <div class="msg-user">
-                <div>
-                    <div class="bubble-user">{safe_content}</div>
-                    <div class="msg-meta" style="text-align:right;">YOU · {msg.get("time","")}</div>
-                </div>
-            </div>"""
-        else:
-            chat_html += f"""
-            <div class="msg-ai">
-                <div>
-                    <div class="bubble-ai">{safe_content}</div>
-                    <div class="msg-meta">⚡ TURBOCHAT · {msg.get("time","")} · {msg.get("tokens","")} tokens · {msg.get("elapsed","")}s</div>
-                </div>
-            </div>"""
-    chat_html += '</div>'
-    st.markdown(chat_html, unsafe_allow_html=True)
-
-# Typing indicator slot sits HERE — above the input box
-typing_slot = st.empty()
-
-# ─── INPUT ROW ────────────────────────────────────────────────────────────────
-st.write("")
-col_input, col_btn = st.columns([6, 0.7])
-with col_input:
-    user_input = st.text_input(
-        "",
-        placeholder="Type your message and press Enter...",
-        label_visibility="collapsed",
-        key=f"user_input_{st.session_state.input_key}"
-    )
-with col_btn:
-    send = st.button("⚡", use_container_width=True)
-
-# ─── GENERATE ─────────────────────────────────────────────────────────────────
-if (send or user_input) and user_input.strip() and user_input.strip() != st.session_state.last_q:
-
-    ts = get_ist_time()
-    st.session_state.messages.append({
-        "role": "user",
-        "content": user_input.strip(),
-        "time": ts
-    })
-    st.session_state.msg_count += 1
-
-    # Show typing indicator above the input
-    typing_slot.markdown("""
-    <div class="msg-ai">
-        <div class="typing-indicator">
-            <div class="typing-dot"></div>
-            <div class="typing-dot"></div>
-            <div class="typing-dot"></div>
+    # Chat history
+    if st.session_state.messages:
+        chat_html = '<div class="chat-wrap">'
+        for msg in st.session_state.messages:
+            ts = msg.get("time", "")
+            if msg["role"] == "user":
+                import html as html_module
+                safe_user = html_module.escape(msg["content"])
+                chat_html += f"""
+                <div class="msg-user">
+                    <div>
+                        <div class="bubble-user">{safe_user}</div>
+                        <div class="msg-meta" style="text-align:right;">YOU · {ts}</div>
+                    </div>
+                </div>"""
+            else:
+                sources_html = ""
+                if msg.get("sources"):
+                    chips = "".join(f'<span class="source-chip">📎 Chunk {i+1}</span>'
+                                    for i, _ in enumerate(msg["sources"]))
+                    sources_html = f"""<div class="sources-wrap">
+                        <div class="source-label">// SOURCE CHUNKS USED</div>
+                        {chips}
+                    </div>"""
+
+                # Escape answer to prevent HTML injection from LLM output
+                import html as html_module
+                safe_answer = html_module.escape(msg["content"]).replace("\n", "<br>")
+
+                chat_html += f"""
+                <div class="msg-ai">
+                    <div>
+                        <div class="bubble-ai">
+                            <div class="answer-label">// QUERYDOCS RESPONSE</div>
+                            <div class="answer-text">{safe_answer}</div>
+                            {sources_html}
+                        </div>
+                        <div class="msg-meta">📚 QUERYDOCS AI · {ts} · {msg.get("elapsed","?")}s</div>
+                    </div>
+                </div>"""
+        chat_html += '</div>'
+        st.markdown(chat_html, unsafe_allow_html=True)
+    else:
+        st.markdown("""
+        <div class="welcome-card" style="padding:20px;">
+            <div style="font-size:1.6rem;margin-bottom:8px;">💬</div>
+            <div class="wc-title" style="font-size:0.85rem;">DOCUMENT LOADED — START ASKING</div>
+            <div class="wc-sub" style="font-size:0.82rem;">Ask anything about the uploaded document.</div>
+            <br>
+            <span class="tip-chip">💡 Summarize this document</span>
+            <span class="tip-chip">💡 What are the key findings?</span>
+            <span class="tip-chip">💡 List all important dates</span>
         </div>
-    </div>
-    """, unsafe_allow_html=True)
-
-    # generate response
-    try:
-        tokenizer, model = load_model()
-
-        prompt = f"<|system|>\n{system_prompt.strip()}\n<|end|>\n"
-        for m in st.session_state.messages:
-            if m["role"] == "user":
-                prompt += f"<|user|>\n{m['content']}\n<|end|>\n"
-            elif m["role"] == "assistant":
-                prompt += f"<|assistant|>\n{m['content']}\n<|end|>\n"
-        prompt += "<|assistant|>\n"
-
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-        if torch.cuda.is_available():
-            input_ids = input_ids.to("cuda")
-
-        start_time = time.time()
-
-        with torch.no_grad():
-            output = model.generate(
-                input_ids,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                repetition_penalty=1.1,
-            )
-
-        elapsed     = round(time.time() - start_time, 1)
-        new_tokens  = output[0][input_ids.shape[-1]:]
-        response    = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
-        token_count = len(new_tokens)
-
-        st.session_state.total_tokens += token_count
-        st.session_state.total_time   += elapsed
-
-        st.session_state.messages.append({
-            "role":    "assistant",
-            "content": response,
-            "time":    get_ist_time(),
-            "tokens":  token_count,
-            "elapsed": elapsed
-        })
-
-    except Exception as e:
+        """, unsafe_allow_html=True)
+
+    # Input row
+    st.markdown('<hr class="glow-div">', unsafe_allow_html=True)
+    col_q, col_btn = st.columns((5, 1))
+
+    with col_q:
+        question = st.text_input(
+            "", placeholder="ask a question about your document...",
+            label_visibility="collapsed", key="question_input"
+        )
+    with col_btn:
+        st.markdown('<div class="send-btn">', unsafe_allow_html=True)
+        ask_btn = st.button("▶ ASK", use_container_width=True)
+        st.markdown('</div>', unsafe_allow_html=True)
+
+    # Generate answer — guard against infinite loop
+    if (ask_btn or question) and question.strip() and question.strip() != st.session_state.last_q:
+        ts = time.strftime("%H:%M")
         st.session_state.messages.append({
-            "role":    "assistant",
-            "content": f"⚠️ Error: {str(e)}",
-            "time":    get_ist_time(),
-            "tokens":  0,
-            "elapsed": 0
+            "role": "user", "content": question.strip(), "time": ts
         })
-
-    typing_slot.empty()
-    st.session_state.last_q   = user_input.strip()
-    st.session_state.input_key += 1
-    st.rerun()
+        st.session_state.q_count += 1
+
+        # typing indicator
+        typing_slot = st.empty()
+        typing_slot.markdown("""
+        <div class="typing-wrap">
+            <div class="typing-box">
+                <div class="t-dot"></div>
+                <div class="t-dot"></div>
+                <div class="t-dot"></div>
+            </div>
+        </div>
+        """, unsafe_allow_html=True)
+
+        try:
+            start   = time.time()
+            answer, sources = get_answer(question.strip(), st.session_state.vectorstore)
+            elapsed = round(time.time() - start, 1)
+
+            st.session_state.messages.append({
+                "role":    "assistant",
+                "content": answer,
+                "sources": sources,
+                "time":    time.strftime("%H:%M"),
+                "elapsed": elapsed
+            })
+        except Exception as e:
+            st.session_state.messages.append({
+                "role": "assistant",
+                "content": f"⚠️ Error generating answer: {str(e)}",
+                "sources": [], "time": time.strftime("%H:%M"), "elapsed": 0
+            })
+
+        typing_slot.empty()
+        st.session_state.last_q = question.strip()
+        st.rerun()