Update app.py

app.py

@@ -1,68 +1,71 @@
 import os
 import streamlit as st
 from groq import Groq
-from langchain.vectorstores import FAISS
 from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
-from langchain.document_loaders import TextLoader
+from langchain.vectorstores import FAISS
 from langchain.text_splitter import CharacterTextSplitter
-from huggingface_hub import hf_hub_download
+from langchain.document_loaders import TextLoader, PyPDFLoader
+from langchain.chains import RetrievalQA
+from langchain.llms.base import LLM
+from tempfile import NamedTemporaryFile
 
-# API key from Hugging Face secrets
+# Load Groq API key from environment variable
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
-# Init Groq client
+# Initialize Groq client
 groq_client = Groq(api_key=GROQ_API_KEY)
 
-# UI setup
-st.set_page_config(page_title="GEO MVP - Generative Engine Optimization", layout="wide")
-st.title("🔍 GEO: Generative Engine Optimization")
+# Define a basic LLM wrapper for Groq
+class GroqLLM(LLM):
+    def __init__(self, model_name="llama3-8b-8192"):
+        self.model_name = model_name
 
-# Upload document
-uploaded_file = st.file_uploader("📄 Upload a .txt file", type=["txt"])
+    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()
 
-if uploaded_file:
-    # Save file
-    with open("data.txt", "wb") as f:
-        f.write(uploaded_file.read())
+    @property
+    def _llm_type(self):
+        return "groq_llm"
 
-    # Load and split
-    loader = TextLoader("data.txt")
-    documents = loader.load()
-    splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    docs = splitter.split_documents(documents)
+# Streamlit UI
+st.title("📚 RAG App with Groq + HuggingFace + Streamlit")
+st.write("Upload a PDF or TXT file, ask a question, and get answers powered by RAG.")
 
-    # Embed
-    st.info("🔎 Generating embeddings...")
-    embeddings = HuggingFaceEmbeddings()
-    vectorstore = FAISS.from_documents(docs, embeddings)
+uploaded_file = st.file_uploader("Upload your document", type=["pdf", "txt"])
 
-    # Build retriever
-    retriever = vectorstore.as_retriever()
+if uploaded_file:
+    with NamedTemporaryFile(delete=False) as tmp_file:
+        tmp_file.write(uploaded_file.read())
+        tmp_path = tmp_file.name
 
-    # Prompt setup
-    prompt_template = PromptTemplate.from_template(
-        "You are an expert assistant. Use the following context to answer accurately:\n\n{context}\n\nQ: {question}\nA:"
-    )
+    # Load document
+    if uploaded_file.type == "application/pdf":
+        loader = PyPDFLoader(tmp_path)
+    else:
+        loader = TextLoader(tmp_path)
 
-    st.success("✅ Data embedded and ready.")
+    docs = loader.load()
 
-    # Query box
-    user_query = st.text_input("💬 Ask a question based on your uploaded file")
+    # Split into chunks
+    splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    texts = splitter.split_documents(docs)
 
-    if user_query:
-        # Retrieve
-        results = retriever.get_relevant_documents(user_query)
-        context = "\n\n".join([doc.page_content for doc in results[:3]])
+    # Create embeddings and FAISS index
+    embeddings = HuggingFaceEmbeddings()
+    db = FAISS.from_documents(texts, embeddings)
 
-        # Call Groq
-        prompt = prompt_template.format(context=context, question=user_query)
-        response = groq_client.chat.completions.create(
-            messages=[{"role": "user", "content": prompt}],
-            model="mixtral-8x7b-32768",  # Or another Groq-supported model
-        )
+    # RAG chain
+    retriever = db.as_retriever()
+    qa_chain = RetrievalQA.from_chain_type(llm=GroqLLM(), retriever=retriever)
+
+    # Input box
+    query = st.text_input("Ask something about the document:")
 
-        answer = response.choices[0].message.content
-        st.markdown("### 📥 Answer")
-        st.write(answer)
+    if query:
+        result = qa_chain.run(query)
+        st.markdown("### 🧠 Answer:")
+        st.success(result)