Update app.py

app.py

@@ -1,19 +1,22 @@
-# Import necessary libraries
 import os
 import fitz  # For PDF extraction
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
-from groq import Groq
+from dotenv import load_dotenv
 import streamlit as st
 
 # Function to extract text from a PDF
 def extract_text_from_pdf(file):
-    doc = fitz.open(stream=file.read(), filetype="pdf")
-    text = ""
-    for page in doc:
-        text += page.get_text()
-    return text
+    try:
+        doc = fitz.open(stream=file.read(), filetype="pdf")
+        text = ""
+        for page in doc:
+            text += page.get_text()
+        return text
+    except Exception as e:
+        st.error(f"Error extracting text: {e}")
+        return ""
 
 # Function to chunk the text
 def chunk_text(text, chunk_size=500):
@@ -35,43 +38,52 @@ embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 
 # Function to generate embeddings
 def generate_embeddings(chunks):
-    embeddings = embedding_model.encode(chunks)
-    return embeddings
+    return embedding_model.encode(chunks)
 
 # Function to store embeddings in FAISS
 def store_embeddings_in_faiss(embeddings):
-    dimension = embeddings.shape[1]
-    index = faiss.IndexFlatL2(dimension)
-    index.add(embeddings)
-    return index
+    try:
+        dimension = embeddings.shape[1]
+        index = faiss.IndexFlatL2(dimension)
+        index.add(embeddings)
+        return index
+    except Exception as e:
+        st.error(f"Error with FAISS: {e}")
+        return None
 
 # Function to retrieve similar chunks
 def retrieve_similar_chunks(query, index, chunks, model):
-    query_embedding = model.encode([query])[0]
-    distances, indices = index.search(np.array([query_embedding]), k=5)
-    return [chunks[i] for i in indices[0]]
-
-# Groq API setup
-# os.environ["GROQ_API_KEY"] = "your_groq_api_key"  # Replace with your API key
-# groq_client = Groq(api_key=os.environ["GROQ_API_KEY"])
-# groq_api_key = os.getenv("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")  # Fetch the API key from environment variables
-from dotenv import load_dotenv
+    try:
+        query_embedding = model.encode([query])[0]
+        distances, indices = index.search(np.array([query_embedding]), k=5)
+        return [chunks[i] for i in indices[0]]
+    except Exception as e:
+        st.error(f"Error retrieving similar chunks: {e}")
+        return []
 
-load_dotenv()  # Load environment variables from .env
-groq_api_key = os.getenv("GROQ_API_KEY")
+# Load environment variables
+load_dotenv()
+groq_api_key = os.getenv("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")
 if not groq_api_key:
-    raise ValueError("The GROQ_API_KEY environment variable is not set.")
+    st.error("The GROQ_API_KEY environment variable is not set.")
+    exit()
+
+# Initialize Groq client
 groq_client = Groq(api_key=groq_api_key)
 
 def query_llm(prompt, model="llama3-8b-8192"):
-    response = groq_client.chat.completions.create(
-        messages=[
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": prompt},
-        ],
-        model=model,
-    )
-    return response.choices[0].message.content
+    try:
+        response = groq_client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": prompt},
+            ],
+            model=model,
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        st.error(f"Error querying LLM: {e}")
+        return "Error in LLM response."
 
 # Streamlit application
 def main():
@@ -80,21 +92,25 @@ def main():
     # File upload
     uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
     if uploaded_file:
-        # Step 1: Extract text from PDF
+        # Extract text
         pdf_text = extract_text_from_pdf(uploaded_file)
+        if not pdf_text:
+            return
+        
         st.write("PDF Text Extracted:")
-        st.write(pdf_text[:500])  # Show a preview of the text
+        st.write(pdf_text[:500])  # Show a preview
         
-        # Step 2: Chunk the text
+        # Chunk the text
         chunks = chunk_text(pdf_text)
-        st.write(f"Text has been split into {len(chunks)} chunks.")
-
-        # Step 3: Generate embeddings and store in FAISS
+        st.write(f"Text split into {len(chunks)} chunks.")
+        
+        # Generate embeddings
         embeddings = np.array(generate_embeddings(chunks))
         index = store_embeddings_in_faiss(embeddings)
-        st.write("Embeddings generated and stored in vector database.")
-
-        # Step 4: User query
+        if index is None:
+            return
+        
+        # Query handling
         query = st.text_input("Enter your query:")
         if query:
             similar_chunks = retrieve_similar_chunks(query, index, chunks, embedding_model)
@@ -102,8 +118,8 @@ def main():
             for i, chunk in enumerate(similar_chunks, start=1):
                 st.write(f"Chunk {i}: {chunk}")
 
-            # Step 5: Query the LLM using Groq API
-            combined_context = " ".join(similar_chunks[:3])  # Combine top 3 chunks
+            # Query the LLM
+            combined_context = " ".join(similar_chunks[:3])
             llm_prompt = f"Context: {combined_context}\n\nQuery: {query}"
             llm_response = query_llm(llm_prompt)
             st.write("LLM Response:")