Update app.py

app.py

@@ -8,6 +8,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 import time
 import base64
+import html as html_module
+from datetime import datetime, timezone, timedelta
+
+# ─── TIMEZONE (IST = UTC+5:30) ────────────────────────────────────────────────
+IST = timezone(timedelta(hours=5, minutes=30))
+def get_ist_time():
+    return datetime.now(IST).strftime("%H:%M")
 
 st.set_page_config(
     page_title="QueryDocs AI",
@@ -72,88 +79,48 @@ footer, #MainMenu { visibility: hidden; }
 }
 .profile-avatar {
     width: 72px; height: 72px;
-    border-radius: 50%;
-    overflow: hidden;
+    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);
+    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%;
-}
+.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;
+    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;
+    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;
+    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;
+    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; }
+.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;
@@ -161,52 +128,22 @@ footer, #MainMenu { visibility: hidden; }
     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: 14px 18px;
-    border-radius: 18px 18px 18px 4px;
-    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);
-}
-.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: #484f58;
-    margin-top: 4px; padding: 0 6px;
-}
-
-.answer-text {
-    white-space: pre-wrap;
-    word-break: break-word;
-    line-height: 1.75;
-    font-size: 0.95rem;
+    font-size: 0.62rem; color: #484f58; margin-top: 4px; padding: 0 6px;
 }
 
 /* ── Source chunks ── */
-.sources-wrap {
-    margin-top: 10px;
-    border-top: 1px solid #21262d;
-    padding-top: 10px;
-}
+.sources-wrap { margin-top: 10px; border-top: 1px solid #21262d; padding-top: 10px; }
 .source-label {
     font-family: 'Share Tech Mono', monospace;
-    font-size: 0.62rem; color: #8b949e;
-    letter-spacing: 2px; margin-bottom: 6px;
+    font-size: 0.62rem; color: #8b949e; letter-spacing: 2px; margin-bottom: 6px;
 }
 .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;
-    margin: 3px 3px; cursor: pointer;
+    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; margin: 3px 3px;
 }
 
 /* ── Typing indicator ── */
@@ -214,33 +151,53 @@ footer, #MainMenu { visibility: hidden; }
 .typing-box {
     display: flex; align-items: center; gap: 5px;
     padding: 14px 18px; background: #161b22;
-    border: 1px solid #30363d;
-    border-radius: 18px 18px 18px 4px;
+    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;
+    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 {
+/* ── Input — merged input+button bar (scoped to main area) ── */
+[data-testid="stAppViewContainer"] div[data-testid="column"]:first-child .stTextInput > div > div > input {
     background: #161b22 !important;
     border: 1px solid #30363d !important;
-    border-radius: 8px !important;
+    border-right: none !important;
+    border-radius: 12px 0 0 12px !important;
     color: #c9d1d9 !important;
     font-family: 'Rajdhani', sans-serif !important;
     font-size: 0.97rem !important;
+    height: 48px !important;
+    padding: 12px 18px !important;
 }
-.stTextInput > div > div > input:focus {
+[data-testid="stAppViewContainer"] div[data-testid="column"]:first-child .stTextInput > div > div > input:focus {
     border-color: #58a6ff !important;
-    box-shadow: 0 0 12px rgba(88,166,255,0.15) !important;
-}
-
-/* ── Buttons ── */
+    box-shadow: none !important;
+    outline: none !important;
+}
+[data-testid="stAppViewContainer"] 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: 1rem !important; letter-spacing: 1px !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;
+}
+[data-testid="stAppViewContainer"] div[data-testid="column"]:last-child .stButton > button:hover {
+    box-shadow: 0 0 24px rgba(31,111,235,0.55) !important;
+    color: #fff !important;
+}
+[data-testid="stAppViewContainer"] div[data-testid="column"]:first-child { padding-right: 0 !important; }
+[data-testid="stAppViewContainer"] div[data-testid="column"]:last-child  { padding-left:  0 !important; }
+
+/* ── Buttons (sidebar/general) ── */
 .stButton > button {
     background: #21262d !important; color: #c9d1d9 !important;
     border: 1px solid #30363d !important; border-radius: 8px !important;
@@ -249,72 +206,30 @@ footer, #MainMenu { visibility: hidden; }
     transition: all 0.25s !important;
 }
 .stButton > button:hover {
-    background: #30363d !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;
+    background: #30363d !important; border-color: #58a6ff !important; color: #58a6ff !important;
 }
 
-/* ── Stats ── */
-.stat-row {
-    display: flex; gap: 10px; flex-wrap: wrap;
-    margin-bottom: 14px;
-}
-.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 ── */
+/* ── Welcome card ── */
 .welcome-card {
     background: #161b22; border: 1px solid #21262d;
-    border-radius: 12px; padding: 30px;
-    text-align: center; margin: 10px 0;
+    border-radius: 12px; padding: 30px; text-align: center; margin: 10px 0;
     animation: fadeInUp 0.6s ease both;
 }
 .wc-icon { font-size: 2.8rem; margin-bottom: 12px; }
-.wc-title {
-    font-family: 'Orbitron', monospace;
-    font-size: 1rem; color: #fff; margin-bottom: 8px;
-}
+.wc-title { font-family: 'Orbitron', monospace; font-size: 1rem; color: #fff; margin-bottom: 8px; }
 .wc-sub { font-size: 0.9rem; color: #8b949e; line-height: 1.65; }
 .tip-chip {
     display: inline-block;
-    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.67rem;
+    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.67rem;
 }
 
 /* ── File uploader ── */
 [data-testid="stFileUploader"] {
-    background: #161b22 !important;
-    border: 2px dashed #30363d !important;
-    border-radius: 10px !important;
-}
-[data-testid="stFileUploader"]:hover {
-    border-color: #58a6ff !important;
+    background: #161b22 !important; border: 2px dashed #30363d !important; border-radius: 10px !important;
 }
+[data-testid="stFileUploader"]:hover { border-color: #58a6ff !important; }
 
 /* ── Glow divider ── */
 .glow-div {
@@ -323,6 +238,26 @@ footer, #MainMenu { visibility: hidden; }
     margin: 16px 0; box-shadow: 0 0 8px rgba(88,166,255,0.3);
 }
 
+/* ── AI bubble native markdown styling ── */
+[data-testid="stAppViewContainer"] [data-testid="column"] p    { color: #c9d1d9; font-size: 0.95rem; line-height: 1.75; margin: 4px 0; }
+[data-testid="stAppViewContainer"] [data-testid="column"] h1,
+[data-testid="stAppViewContainer"] [data-testid="column"] h2,
+[data-testid="stAppViewContainer"] [data-testid="column"] h3   { color: #fff; margin: 10px 0 4px; }
+[data-testid="stAppViewContainer"] [data-testid="column"] ul,
+[data-testid="stAppViewContainer"] [data-testid="column"] ol   { color: #c9d1d9; padding-left: 20px; }
+[data-testid="stAppViewContainer"] [data-testid="column"] code { background: #0d1117; color: #58a6ff; padding: 2px 6px; border-radius: 4px; font-size: 0.85rem; }
+[data-testid="stAppViewContainer"] [data-testid="column"] pre  { background: #0d1117; border: 1px solid #30363d; border-radius: 8px; padding: 12px; }
+[data-testid="stAppViewContainer"] [data-testid="column"] strong { color: #fff; }
+
+/* ── Stats ── */
+.stat-row { display: flex; gap: 10px; flex-wrap: wrap; margin-bottom: 14px; }
+.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; }
+
 /* ── 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)} }
@@ -337,20 +272,20 @@ footer, #MainMenu { visibility: hidden; }
 """, unsafe_allow_html=True)
 
 # ─── 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      = ""
+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         = ""
+if "input_key"      not in st.session_state: st.session_state.input_key      = 0   
+if "pending_input"  not in st.session_state: st.session_state.pending_input  = ""  
 
 # ─── MODEL LOADERS ────────────────────────────────────────────────────────────
 @st.cache_resource(show_spinner=False)
 def load_embeddings():
-    return HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2"
-    )
+    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 @st.cache_resource(show_spinner=False)
 def load_llm():
@@ -383,8 +318,8 @@ def process_pdf(uploaded_file):
     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)
+    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)
@@ -413,7 +348,7 @@ QUESTION:
     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:
@@ -441,13 +376,9 @@ with st.sidebar:
     </div>
     """, unsafe_allow_html=True)
 
-    # 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)
 
-    uploaded_file = st.file_uploader(
-        "Upload PDF", type=["pdf"],
-        label_visibility="collapsed"
-    )
+    uploaded_file = st.file_uploader("Upload PDF", type=["pdf"], label_visibility="collapsed")
 
     if uploaded_file:
         if st.session_state.pdf_name != uploaded_file.name:
@@ -466,7 +397,6 @@ with st.sidebar:
 
     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>
@@ -492,13 +422,11 @@ with st.sidebar:
         """, unsafe_allow_html=True)
         st.markdown("---")
 
-    # Clear button
     if st.button("🗑️ CLEAR CHAT", use_container_width=True):
         st.session_state.messages = []
         st.session_state.q_count  = 0
         st.rerun()
 
-    # New PDF button
     if st.button("📄 LOAD NEW PDF", use_container_width=True):
         st.session_state.vectorstore = None
         st.session_state.pdf_name    = None
@@ -509,23 +437,22 @@ with st.sidebar:
         st.rerun()
 
 # ─── MAIN AREA ────────────────────────────────────────────────────────────────
-# Header
 st.markdown("""
 <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 class="app-sub">INTELLIGENT DOCUMENT Q&amp;A · RAG PIPELINE</div>
     </div>
 </div>
 """, unsafe_allow_html=True)
 
-# PDF banner (when loaded)
 if st.session_state.pdf_name:
+    safe_pdf_name = html_module.escape(st.session_state.pdf_name)   # FIX: XSS
     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-name">{safe_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);
@@ -534,7 +461,6 @@ if st.session_state.pdf_name:
     </div>
     """, unsafe_allow_html=True)
 
-# Chat area
 if not st.session_state.vectorstore:
     st.markdown("""
     <div class="welcome-card">
@@ -542,7 +468,7 @@ if not st.session_state.vectorstore:
         <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.
+            Powered by RAG pipeline — context-aware answers.
         </div>
         <br>
         <span class="tip-chip">📋 Legal documents</span>
@@ -553,48 +479,57 @@ if not st.session_state.vectorstore:
     """, unsafe_allow_html=True)
 
 else:
-    # Chat history
+    # ── 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"""
+                safe_user = html_module.escape(msg["content"])   # FIX: XSS
+                st.markdown(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>"""
+                </div>""", unsafe_allow_html=True)
             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}
+                
+                col_ai, col_space = st.columns([4, 1])
+                with col_ai:
+                   
+                    st.markdown(f"""
+                    <div style="background:#161b22;color:#c9d1d9;padding:14px 18px 6px 18px;
+                        border-radius:18px 18px 0 4px;font-size:0.95rem;line-height:1.75;
+                        border:1px solid #30363d;border-bottom:none;
+                        box-shadow:0 4px 15px rgba(0,0,0,0.3);">
+                        <div style="font-family:'Share Tech Mono',monospace;font-size:0.65rem;
+                                    color:#58a6ff;letter-spacing:2px;margin-bottom:6px;">// QUERYDOCS RESPONSE</div>
+                    </div>""", unsafe_allow_html=True)
+                    
+                    st.markdown(f'<div style="background:#161b22;padding:0 18px 6px 18px;border-left:1px solid #30363d;border-right:1px solid #30363d;">', unsafe_allow_html=True)
+                    st.markdown(msg["content"])
+                    st.markdown('</div>', unsafe_allow_html=True)
+                    # Source chips
+                    if msg.get("sources"):
+                        chips = "".join(f'<span class="source-chip">📎 Chunk {i+1}</span>'
+                                        for i, _ in enumerate(msg["sources"]))
+                        st.markdown(f"""
+                        <div style="background:#161b22;padding:6px 18px 8px 18px;
+                            border-left:1px solid #30363d;border-right:1px solid #30363d;">
+                            <div class="sources-wrap">
+                                <div class="source-label">// SOURCE CHUNKS USED</div>
+                                {chips}
+                            </div>
+                        </div>""", unsafe_allow_html=True)
+                    # Bubble footer
+                    st.markdown(f"""
+                    <div style="background:#161b22;padding:6px 18px 12px 18px;
+                        border-radius:0 0 18px 4px;border:1px solid #30363d;border-top:none;
+                        box-shadow:0 4px 15px rgba(0,0,0,0.3);margin-bottom:4px;">
+                        <div style="font-family:'Share Tech Mono',monospace;font-size:0.62rem;color:#484f58;">
+                            📚 QUERYDOCS AI · {ts} · {msg.get("elapsed","?")}s
                         </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)
+                    </div>""", unsafe_allow_html=True)
     else:
         st.markdown("""
         <div class="welcome-card" style="padding:20px;">
@@ -608,30 +543,44 @@ else:
         </div>
         """, unsafe_allow_html=True)
 
-    # Input row
+    typing_slot = st.empty()
+
+    # ── Input row — merged input+button ──
     st.markdown('<hr class="glow-div">', unsafe_allow_html=True)
-    col_q, col_btn = st.columns((5, 1))
+    _current_key = f"question_input_{st.session_state.input_key}"
+
+    def _sync_pending():
+        st.session_state.pending_input = st.session_state.get(_current_key, "")
 
+    col_q, col_btn = st.columns([6, 0.7])
     with col_q:
         question = st.text_input(
             "", placeholder="ask a question about your document...",
-            label_visibility="collapsed", key="question_input"
+            label_visibility="collapsed",
+            key=_current_key,
+            on_change=_sync_pending
         )
+        if question:
+            st.session_state.pending_input = question
     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")
+    trigger_q = (st.session_state.get("pending_input", "") or question).strip()
+
+    # ── Generate answer ──
+    if trigger_q and trigger_q != st.session_state.last_q:
+
+        # FIX: set last_q and clear input IMMEDIATELY to prevent double-trigger
+        st.session_state.last_q        = trigger_q
+        st.session_state.input_key    += 1
+        st.session_state.pending_input = ""
+
+        ts = get_ist_time()   # FIX: IST time instead of UTC server time
         st.session_state.messages.append({
-            "role": "user", "content": question.strip(), "time": ts
+            "role": "user", "content": trigger_q, "time": ts
         })
         st.session_state.q_count += 1
 
-        # typing indicator
-        typing_slot = st.empty()
         typing_slot.markdown("""
         <div class="typing-wrap">
             <div class="typing-box">
@@ -643,24 +592,25 @@ else:
         """, unsafe_allow_html=True)
 
         try:
-            start   = time.time()
-            answer, sources = get_answer(question.strip(), st.session_state.vectorstore)
-            elapsed = round(time.time() - start, 1)
+            start          = time.time()
+            answer, sources = get_answer(trigger_q, 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"),
+                "time":    get_ist_time(),   # FIX: IST time
                 "elapsed": elapsed
             })
         except Exception as e:
             st.session_state.messages.append({
-                "role": "assistant",
+                "role":    "assistant",
                 "content": f"⚠️ Error generating answer: {str(e)}",
-                "sources": [], "time": time.strftime("%H:%M"), "elapsed": 0
+                "sources": [],
+                "time":    get_ist_time(),
+                "elapsed": 0
             })
 
         typing_slot.empty()
-        st.session_state.last_q = question.strip()
         st.rerun()