Update app.py

app.py

@@ -1,78 +1,60 @@
 import os
 import streamlit as st
+from langchain.embeddings import Embedding
 from groq import Groq
-from langchain.embeddings import BaseEmbedding
+from langchain.chains import RetrievalQA
 from langchain.vectorstores import FAISS
 from langchain.document_loaders import PyPDFLoader
-from langchain.chains.question_answering import load_qa_chain
-from langchain.llms import OpenAI  # Keep this if you're still using OpenAI for QA model, otherwise replace it
-from langchain.prompts import PromptTemplate
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.llms import OpenAI
+from langchain.agents import initialize_agent
+from langchain.agents import Tool
+
+# Set up Groq API
+groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+
+# Define a custom embedding class for Groq
+class GroqEmbedding(Embedding):
+    def __init__(self, model="groq-embedding-model", api_key=None):
+        self.model = model
+        self.client = Groq(api_key=api_key or os.getenv("GROQ_API_KEY"))
+    
+    def embed_documents(self, texts):
+        # Use Groq's API to generate embeddings
+        embeddings = self.client.embed_documents(texts, model=self.model)
+        return embeddings
+    
+    def embed_query(self, query):
+        # Use Groq's API to generate query embedding
+        return self.client.embed_query(query, model=self.model)
 
-# Groq API setup for embeddings
-class GroqEmbedding(BaseEmbedding):
-    def __init__(self, api_key: str):
-        self.client = Groq(api_key=api_key)
+# Streamlit App UI
+st.title("PDF Question-Answering with Groq Embeddings")
 
-    def embed_documents(self, texts: list) -> list:
-        embeddings = []
-        for text in texts:
-            response = self.client.embeddings.create(input=text)
-            embeddings.append(response['data'])
-        return embeddings
+uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
 
-# Load documents from uploaded PDF file
-def load_documents(uploaded_file):
+# Process the uploaded PDF
+if uploaded_file is not None:
     loader = PyPDFLoader(uploaded_file)
     documents = loader.load()
-    return documents
-
-# Create FAISS vector database with Groq embeddings
-def create_vector_db(documents):
-    # Use Groq embeddings
-    embeddings = GroqEmbedding(api_key="your-groq-api-key")  # Pass your Groq API key
-    vector_db = FAISS.from_documents(documents, embeddings)
-    return vector_db
-
-# Function to perform QA with the uploaded documents
-def perform_qa(vector_db, query):
-    # Set up the prompt and model for QA
-    prompt_template = "Answer the following question based on the documents: {question}"
-    prompt = PromptTemplate(input_variables=["question"], template=prompt_template)
-    qa_chain = load_qa_chain(OpenAI(), chain_type="stuff", prompt=prompt)  # Keep OpenAI model for QA
-
-    # Query the vector DB to retrieve the most relevant documents
-    results = vector_db.similarity_search(query)
-
-    # Perform QA using the chain
-    answer = qa_chain.run(input_documents=results, question=query)
-    return answer
-
-# Streamlit UI setup
-def main():
-    st.title("Document Upload and Question Answering")
-
-    # Upload PDF file
-    uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])
-    if uploaded_file:
-        st.write("File uploaded successfully!")
-
-        try:
-            # Load documents from the uploaded PDF
-            documents = load_documents(uploaded_file)
-
-            # Create a vector DB using Groq embeddings
-            vector_db = create_vector_db(documents)
-
-            # User query for Q&A
-            query = st.text_input("Ask a question based on the uploaded document:")
-
-            if query:
-                # Get the answer for the query
-                answer = perform_qa(vector_db, query)
-                st.write("Answer:", answer)
-
-        except Exception as e:
-            st.error(f"Error loading client or processing query: {e}")
+    
+    # Split documents into smaller chunks for better processing
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    split_docs = text_splitter.split_documents(documents)
+
+    # Create embeddings using Groq
+    embeddings = GroqEmbedding(api_key=os.getenv("GROQ_API_KEY"))
+    
+    # Create a FAISS vector store
+    vector_db = FAISS.from_documents(split_docs, embeddings)
+    
+    # Initialize the retrieval-based QA system
+    qa = RetrievalQA.from_chain_type(llm=OpenAI(), chain_type="stuff", vectorstore=vector_db)
+    
+    # User input for querying the PDF content
+    query = st.text_input("Ask a question about the PDF:")
+
+    if query:
+        result = qa.run(query)
+        st.write("Answer:", result)
 
-if __name__ == "__main__":
-    main()