Update app.py

app.py

@@ -10,234 +10,282 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 import torch
 
-# ---------------- CONFIG ----------------
+# ---------------- CONFIGURATION ----------------
 logging.basicConfig(level=logging.INFO)
 
+# Load API key from Hugging Face secrets
 GROQ_API_KEY = st.secrets.get("GROQ_API_KEY", os.environ.get("GROQ_API_KEY"))
 GROQ_MODEL = "llama-3.1-8b-instant"
 
+# Initialize Groq client
 client = None
 if GROQ_API_KEY:
     try:
         client = Groq(api_key=GROQ_API_KEY)
+        st.success("✅ Groq client initialized successfully.")
     except Exception as e:
-        logging.error(e)
+        st.error(f"❌ Failed to initialize Groq client: {e}")
+        client = None
+else:
+    st.warning("⚠️ GROQ_API_KEY not found. Please add it to Hugging Face secrets.")
 
-# ---------------- PAGE ----------------
+# ---------------- STREAMLIT UI SETUP ----------------
 st.set_page_config(page_title="PDF Assistant", page_icon="📘", layout="wide")
 
-# ---------------- CSS ----------------
+# ---------------- CSS (Your exact UI) ----------------
 st.markdown("""
 <style>
 :root {
-    --primary:#1e3a8a;
-    --bg:#0e1117;
-    --bg2:#1a1d29;
-    --text:#f0f2f6;
-}
-
-/* FIX SIDEBAR */
-section[data-testid="stSidebar"] {
-    position: fixed;
-    height: 100vh;
-    overflow-y: auto;
-}
-section[data-testid="stSidebar"] > div {
-    padding-top: 0.5rem !important;
-}
-
-/* MAIN OFFSET */
-.main {
-    margin-left: 300px;
-}
-
-/* CHAT AREA */
-.chat-area {
-    height: calc(100vh - 230px);
-    overflow-y: auto;
-    padding: 1rem 2rem;
+    --primary-color: #1e3a8a;
+    --background-color: #0e1117;
+    --secondary-background-color: #1a1d29;
+    --text-color: #f0f2f6;
 }
 
 .chat-user {
-    background:#2d3748;
-    padding:14px;
-    border-radius:18px 18px 4px 18px;
-    margin:12px 0 12px auto;
-    max-width:80%;
+    background: #2d3748;
+    padding: 12px;
+    border-radius: 10px 10px 2px 10px;
+    margin: 6px 0 6px auto;
+    max-width: 85%;
+    text-align: right;
+    color: var(--text-color);
 }
-
 .chat-bot {
-    background:var(--primary);
-    padding:14px;
-    border-radius:18px 18px 18px 4px;
-    margin:12px auto 12px 0;
-    max-width:80%;
-    color:white;
+    background: var(--primary-color);
+    padding: 12px;
+    border-radius: 10px 10px 10px 2px;
+    margin: 6px auto 6px 0;
+    max-width: 85%;
+    text-align: left;
+    color: #ffffff;
 }
 
 .sources {
-    font-size:0.75em;
-    opacity:0.7;
-    margin-top:8px;
-    border-top:1px solid rgba(255,255,255,0.15);
-    padding-top:6px;
+    font-size: 0.8em;
+    opacity: 0.7;
+    margin-top: 10px;
+    border-top: 1px solid rgba(255, 255, 255, 0.1);
+    padding-top: 5px;
 }
 
-/* INPUT BAR */
-.input-bar {
-    position: sticky;
+.footer {
+    position: fixed;
+    left: 0;
     bottom: 0;
-    background: var(--bg);
-    padding: 1rem 2rem;
-    border-top:1px solid rgba(255,255,255,0.15);
+    width: 100%;
+    background-color: var(--secondary-background-color);
+    color: var(--text-color);
+    text-align: center;
+    padding: 10px;
+    font-size: 0.85em;
+    border-top: 1px solid rgba(255, 255, 255, 0.1);
 }
-
-.input-box {
-    max-width:800px;
-    margin:auto;
+.footer a {
+    color: var(--primary-color);
+    text-decoration: none;
+    font-weight: bold;
+}
+.footer a:hover {
+    text-decoration: underline;
 }
 </style>
 """, unsafe_allow_html=True)
 
-# ---------------- STATE ----------------
+# ---------------- SESSION STATE ----------------
 if "chat" not in st.session_state:
     st.session_state.chat = []
+
 if "vectorstore" not in st.session_state:
     st.session_state.vectorstore = None
+
 if "retriever" not in st.session_state:
     st.session_state.retriever = None
-if "pdf" not in st.session_state:
-    st.session_state.pdf = None
+
+if "uploaded_file_name" not in st.session_state:
+    st.session_state.uploaded_file_name = None
+
 if "uploader_key" not in st.session_state:
     st.session_state.uploader_key = 0
 
 # ---------------- FUNCTIONS ----------------
-def clear_all():
+def clear_chat_history():
     st.session_state.chat = []
+
+def clear_memory():
     st.session_state.vectorstore = None
     st.session_state.retriever = None
-    st.session_state.pdf = None
+    st.session_state.uploaded_file_name = None
     st.session_state.uploader_key += 1
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    st.success("Memory cleared")
-
-def process_pdf(uploaded):
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
-        tmp.write(uploaded.getvalue())
-        path = tmp.name
-
-    loader = PyPDFLoader(path)
-    docs = loader.load()
-
-    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=50)
-    chunks = splitter.split_documents(docs)
-
-    embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={"device":"cpu"},
-        encode_kwargs={"normalize_embeddings":True}
-    )
-
-    st.session_state.vectorstore = Chroma.from_documents(chunks, embeddings)
-    st.session_state.retriever = st.session_state.vectorstore.as_retriever(k=3)
-
-    os.unlink(path)
-    return len(chunks)
+    st.success("Memory cleared. Please upload a new PDF.")
 
-def ask_llm(q):
-    docs = st.session_state.retriever.invoke(q)
-    ctx = "\n\n".join(d.page_content for d in docs)
-
-    prompt = f"""
-Use ONLY the context. If answer not found, say:
+def process_pdf(uploaded_file):
+    """Process uploaded PDF and create vectorstore."""
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+            tmp.write(uploaded_file.getvalue())
+            path = tmp.name
+        
+        # Load PDF
+        loader = PyPDFLoader(path)
+        docs = loader.load()
+        
+        # Split into chunks
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=800,
+            chunk_overlap=50
+        )
+        chunks = splitter.split_documents(docs)
+        
+        # Create embeddings
+        embeddings = HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-MiniLM-L6-v2",
+            model_kwargs={"device": "cpu"},
+            encode_kwargs={"normalize_embeddings": True}
+        )
+        
+        # Create vectorstore
+        vectorstore = Chroma.from_documents(chunks, embeddings)
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+        
+        # Store in session state
+        st.session_state.vectorstore = vectorstore
+        st.session_state.retriever = retriever
+        
+        # Cleanup
+        if os.path.exists(path):
+            os.unlink(path)
+        
+        return len(chunks)
+        
+    except Exception as e:
+        st.error(f"Error processing PDF: {str(e)}")
+        return None
+
+def ask_question(question):
+    """Retrieve and generate answer for the question."""
+    if not client:
+        return None, 0, "Groq client is not initialized. Check API key setup."
+    
+    if not st.session_state.retriever:
+        return None, 0, "Upload PDF first to initialize the knowledge base."
+    
+    try:
+        # Retrieve relevant chunks
+        docs = st.session_state.retriever.invoke(question)
+        context = "\n\n".join(d.page_content for d in docs)
+        
+        # Build prompt
+        prompt = f"""
+You are a strict RAG Q&A assistant.
+Use ONLY the context provided. If the answer is not found, reply:
 "I cannot find this in the PDF."
 
-CONTEXT:
-{ctx}
+---------------- CONTEXT ----------------
+{context}
+-----------------------------------------
 
-QUESTION:
-{q}
+QUESTION: {question}
+
+FINAL ANSWER:
 """
+        
+        # Call Groq API
+        response = client.chat.completions.create(
+            model=GROQ_MODEL,
+            messages=[
+                {"role": "system",
+                 "content": "Use only the PDF content. If answer not found, say: 'I cannot find this in the PDF.'"},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.0
+        )
+        
+        answer = response.choices[0].message.content.strip()
+        return answer, len(docs), None
+        
+    except APIError as e:
+        return None, 0, f"Groq API Error: {str(e)}"
+    except Exception as e:
+        return None, 0, f"General error: {str(e)}"
 
-    r = client.chat.completions.create(
-        model=GROQ_MODEL,
-        messages=[
-            {"role":"system","content":"Answer strictly from PDF"},
-            {"role":"user","content":prompt}
-        ],
-        temperature=0.0
-    )
-    return r.choices[0].message.content.strip(), len(docs)
-
-# ================= SIDEBAR =================
+# ---------------- UI COMPONENTS ----------------
+st.title("📘 PDF Assistant")
+
+# Sidebar Controls
 with st.sidebar:
-    st.markdown("""
-    <h1 style="text-align:center;margin-bottom:0.2rem;">📘 PDF Assistant</h1>
-    <div style="text-align:center;font-size:0.85em;margin-bottom:1rem;">
-        Created by <a href="https://www.linkedin.com/in/abhishek-iitr/" target="_blank">Abhishek Saxena</a>
-    </div>
-    """, unsafe_allow_html=True)
-
-    st.subheader("Upload PDF")
-    uploaded = st.file_uploader(
-        "PDF",
-        type=["pdf"],
-        key=st.session_state.uploader_key,
-        label_visibility="collapsed"
-    )
-
-    if uploaded and uploaded.name != st.session_state.pdf:
-        st.session_state.chat = []
-        with st.spinner("Processing PDF..."):
-            chunks = process_pdf(uploaded)
-            st.session_state.pdf = uploaded.name
-            st.success(f"PDF processed ({chunks} chunks)")
-        st.rerun()
-
-    st.subheader("Controls")
-    st.button("🧹 Clear Memory", on_click=clear_all, use_container_width=True)
-
-    st.subheader("Status")
-    if st.session_state.pdf:
-        st.success(f"Active PDF:\n`{st.session_state.pdf}`")
+    st.header("Controls")
+    st.button("🗑️ Clear Chat History", on_click=clear_chat_history, use_container_width=True)
+    st.button("🔥 Clear PDF Memory", on_click=clear_memory, use_container_width=True)
+    
+    st.markdown("---")
+    if st.session_state.uploaded_file_name:
+        st.success(f"✅ **Active PDF:**\n `{st.session_state.uploaded_file_name}`")
     else:
-        st.info("Upload a PDF to start")
+        st.warning("⬆️ Upload a PDF to start chatting!")
 
-# ================= CHAT =================
-st.markdown('<div class="chat-area">', unsafe_allow_html=True)
+# File Upload
+uploaded = st.file_uploader(
+    "Upload your PDF",
+    type=["pdf"],
+    key=st.session_state.uploader_key
+)
 
-if not st.session_state.chat:
-    st.info("Ask a question about your document to begin.")
-else:
-    for role, msg in st.session_state.chat:
-        if role == "user":
-            st.markdown(f"<div class='chat-user'>{msg}</div>", unsafe_allow_html=True)
+if uploaded and uploaded.name != st.session_state.uploaded_file_name:
+    st.session_state.uploaded_file_name = None
+    st.session_state.chat = []
+    
+    with st.spinner(f"Processing '{uploaded.name}'..."):
+        chunks_count = process_pdf(uploaded)
+        
+        if chunks_count is not None:
+            st.success(f"✅ PDF processed successfully! {chunks_count} chunks created.")
+            st.session_state.uploaded_file_name = uploaded.name
         else:
-            st.markdown(f"<div class='chat-bot'>{msg}</div>", unsafe_allow_html=True)
-
-st.markdown('</div>', unsafe_allow_html=True)
-
-# ================= INPUT =================
-st.markdown('<div class="input-bar"><div class="input-box">', unsafe_allow_html=True)
-
-with st.form("chat", clear_on_submit=True):
-    q = st.text_input(
-        "Ask",
-        disabled=st.session_state.pdf is None,
-        placeholder="Ask anything about your document…",
-        label_visibility="collapsed"
-    )
-    send = st.form_submit_button("Send")
-
-st.markdown('</div></div>', unsafe_allow_html=True)
+            st.error("❌ Failed to process PDF")
+            st.session_state.uploaded_file_name = None
+    
+    st.rerun()
 
-if send and q:
-    st.session_state.chat.append(("user", q))
+# Chat Input
+disabled_input = st.session_state.uploaded_file_name is None or client is None
+question = st.text_input(
+    "Ask a question about the loaded PDF:",
+    key="question_input",
+    disabled=disabled_input
+)
+
+if st.button("Send", disabled=disabled_input) and question:
+    # Add user query to chat history
+    st.session_state.chat.append(("user", question))
+    
+    # Get answer
     with st.spinner("Thinking..."):
-        ans, src = ask_llm(q)
-        st.session_state.chat.append(
-            ("bot", f"{ans}<div class='sources'>Chunks used: {src}</div>")
-        )
+        answer, sources, error = ask_question(question)
+        
+        if answer:
+            bot_message = f"{answer}<div class='sources'>Context Chunks Used: {sources}</div>"
+            st.session_state.chat.append(("bot", bot_message))
+        else:
+            st.session_state.chat.append(("bot", f"🔴 **Error:** {error}"))
+    
     st.rerun()
+
+# Display Chat History
+st.markdown("## Chat History")
+for role, msg in st.session_state.chat:
+    if role == "user":
+        st.markdown(f"<div class='chat-user'>{msg}</div>", unsafe_allow_html=True)
+    else:
+        st.markdown(f"<div class='chat-bot'>{msg}</div>", unsafe_allow_html=True)
+
+# Footer
+footer_html = """
+<div class="footer">
+    Created by <a href="https://www.linkedin.com/in/abhishek-iitr/" target="_blank">Abhishek Saxena</a>
+</div>
+"""
+st.markdown(footer_html, unsafe_allow_html=True)