Update app.py

app.py

@@ -1,76 +1,76 @@
 import os
+import tempfile
 import streamlit as st
-from groq import Groq
-from langchain_community.embeddings import HuggingFaceEmbeddings
+
+from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.vectorstores import FAISS
-from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.document_loaders import TextLoader, PyPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
-from langchain.llms.base import LLM
-from tempfile import NamedTemporaryFile
+from langchain.prompts import PromptTemplate
+from langchain_groq import GroqLLM
 
-# Load Groq API key from environment variable
+# Set environment variables (You can also use os.environ or Streamlit secrets)
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-if not GROQ_API_KEY:
-    st.error("❌ GROQ_API_KEY is not set in environment variables.")
-    st.stop()
-
-# Initialize Groq client
-groq_client = Groq(api_key=GROQ_API_KEY)
-
-# Define a basic LLM wrapper for Groq
-class GroqLLM(LLM):
-    def __init__(self, model_name="llama3-8b-8192"):
-        self.model_name = model_name
-
-    def _call(self, prompt, stop=None):
-        response = groq_client.chat.completions.create(
-            model=self.model_name,
-            messages=[{"role": "user", "content": prompt}],
-        )
-        return response.choices[0].message.content.strip()
-
-    @property
-    def _llm_type(self):
-        return "groq_llm"
-
-# Streamlit UI
-st.set_page_config(page_title="Groq RAG App", layout="centered")
-st.title("📚 RAG App with Groq + LangChain + FAISS")
-st.write("Upload a PDF or TXT file, ask a question, and get smart answers.")
-
-uploaded_file = st.file_uploader("Upload your document", type=["pdf", "txt"])
-
-if uploaded_file:
-    with NamedTemporaryFile(delete=False) as tmp_file:
+HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
+
+# Initialize Groq LLM
+llm = GroqLLM(
+    api_key=GROQ_API_KEY,
+    model="llama3-8b-8192",  # <- correct param
+    temperature=0.1
+)
+
+# HuggingFace Embeddings
+embedding = HuggingFaceEmbeddings()
+
+st.title("📄 RAG Chat with Groq + HuggingFace")
+
+# Upload PDF
+uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
+
+user_query = st.text_input("Ask something about the document")
+submit_button = st.button("Submit")
+
+if uploaded_file and submit_button:
+    # Save PDF temporarily
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
         tmp_file.write(uploaded_file.read())
         tmp_path = tmp_file.name
 
-    # Load document based on file type
-    if uploaded_file.type == "application/pdf":
-        loader = PyPDFLoader(tmp_path)
-    else:
-        loader = TextLoader(tmp_path)
+    # Load and split
+    loader = PyPDFLoader(tmp_path)
+    pages = loader.load_and_split()
+
+    # Create FAISS vectorstore
+    vectorstore = FAISS.from_documents(pages, embedding)
+    retriever = vectorstore.as_retriever()
 
-    docs = loader.load()
+    # Custom prompt (optional)
+    prompt_template = PromptTemplate(
+        input_variables=["context", "question"],
+        template="""
+        Use the following context to answer the question. Be concise and accurate.
 
-    # Split into chunks
-    splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    texts = splitter.split_documents(docs)
+        Context: {context}
 
-    # Create embeddings and FAISS index
-    embeddings = HuggingFaceEmbeddings()
-    db = FAISS.from_documents(texts, embeddings)
+        Question: {question}
+        """
+    )
 
-    # Set up RAG chain
-    retriever = db.as_retriever()
-    qa_chain = RetrievalQA.from_chain_type(llm=GroqLLM(), retriever=retriever)
+    # Create QA chain
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=True,
+        chain_type_kwargs={"prompt": prompt_template}
+    )
 
-    # Input box for user question
-    query = st.text_input("🔍 Ask a question about the document:")
+    # Run QA
+    result = qa_chain({"query": user_query})
+    st.markdown("### 💬 Answer")
+    st.write(result["result"])
 
-    if query:
-        with st.spinner("Thinking..."):
-            result = qa_chain.run(query)
-        st.markdown("### 🧠 Answer:")
-        st.success(result)
+    # Optional: Show sources
+    with st.expander("📄 Sources"):
+        for doc in result["source_documents"]:
+            st.write(doc.metadata["source"])