Update app.py

app.py

@@ -6,9 +6,23 @@ from langchain.text_splitter import CharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
+from langchain.llms.base import LLM
 from PyPDF2 import PdfReader
 import requests
 
+# Custom LLM wrapper for Groq
+class GroqLLM(LLM):
+    def _call(self, prompt: str, stop: list = None) -> str:
+        response = client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model="llama-3.3-70b-versatile",
+        )
+        return response.choices[0].message.content
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom_groq"
+
 # Initialize Groq client
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
@@ -34,17 +48,14 @@ for page in pdf_reader.pages:
 text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 chunks = text_splitter.split_text(text)
 
-# Use HuggingFace embeddings (Groq-compatible)
+# Use HuggingFace embeddings
 embeddings = HuggingFaceEmbeddings()
 vectorstore = FAISS.from_texts(chunks, embeddings)
 
 # Set up retrieval-based QA
 retriever = vectorstore.as_retriever()
 qa_chain = RetrievalQA.from_chain_type(
-    llm=lambda query: client.chat.completions.create(
-        messages=[{"role": "user", "content": query}],
-        model="llama-3.3-70b-versatile",
-    ).choices[0].message.content,
+    llm=GroqLLM(),  # Pass the custom Groq LLM here
     retriever=retriever,
     return_source_documents=True,
 )