Create components/pdf_chat.py

components/pdf_chat.py

@@ -1,37 +1,47 @@
-# pdf_chat.py
+# components/pdf_chat.py
 
 import streamlit as st
-import os
-from utils.retriever import load_vectorstore
-from langchain_groq import ChatGroq
 from langchain_core.messages import AIMessage
+from utils.retriever import load_vectorstore
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.chat_models import ChatGroq
+
 
 def run_pdf_qa():
-    st.title("💬 Ask Your PDF")
+    st.header("💬 Ask Your PDF")
 
     if "pdf_path" not in st.session_state:
-        st.warning("⚠ Please upload a PDF first from the '📤 Upload Business Docs' section.")
+        st.warning("⚠️ Please upload a PDF first in the 'Upload Business Docs' section.")
         return
 
-    # Show current file
-    st.markdown(f"**Using file:** `{os.path.basename(st.session_state.pdf_path)}`")
+    st.markdown(f"**Using file:** `{st.session_state.pdf_path.split('/')[-1]}`")
 
-    question = st.text_input("❓ Ask a question based on your uploaded PDF")
+    # User Question Input
+    question = st.text_input("❓ Ask a question based on your uploaded PDF", placeholder="e.g. What is the business model?")
 
     if question:
-        with st.spinner("🔍 Thinking..."):
-            retriever = load_vectorstore(st.session_state.pdf_path)
-            docs = retriever.get_relevant_documents(question)
-            context = "\n\n".join([doc.page_content for doc in docs])
-
-            chain = ChatGroq(model="llama3-8b-8192", temperature=0.3)
-            response = chain.invoke(
-                f"""Answer the question based on the context below:\n\n{context}\n\nQuestion: {question}"""
-            )
-
-            st.markdown("### 💡 Answer:")
-            st.markdown(response.content if isinstance(response, AIMessage) else str(response))
-
-            with st.expander("📄 Source Context"):
-                for i, doc in enumerate(docs):
-                    st.markdown(f"**Snippet {i+1}**: {doc.page_content[:500]}...")
+        # Load vectorstore retriever
+        retriever = load_vectorstore(st.session_state.pdf_path)
+
+        # Prompt and Model
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", "You are a helpful assistant. Use the retrieved context to answer the user's question."),
+            ("human", "Context:\n{context}\n\nQuestion: {question}")
+        ])
+        model = ChatGroq(temperature=0.2, model_name="LLaMA3-8b-8192")
+
+        rag_chain = (
+            {"context": retriever | RunnablePassthrough(), "question": RunnablePassthrough()}
+            | prompt
+            | model
+            | StrOutputParser()
+        )
+
+        with st.spinner("Thinking..."):
+            response = rag_chain.invoke(question)
+
+        # Display Answer
+        st.markdown("### 💡 Answer:")
+        st.write(response)