Update pdf_bot.py

pdf_bot.py

@@ -5,28 +5,53 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import RetrievalQA
+from langchain_core.language_models import BaseChatModel
+from langchain_core.outputs import ChatResult, ChatGeneration
+from langchain_core.messages import AIMessage
 from groq import Groq
 
 load_dotenv()
 groq_api_key = os.getenv("GROQ_API_KEY")
-hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-# Load PDF and prepare QA chain
+# ✅ Custom wrapper
+class ChatGroq(BaseChatModel):
+    def __init__(self, model="llama3-8b-8192", temperature=0.3, api_key=None):
+        self.client = Groq(api_key=api_key)
+        self.model = model
+        self.temperature = temperature
+
+    def _generate(self, messages, stop=None):
+        prompt = [{"role": "user", "content": self._get_message_text(messages)}]
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=prompt,
+            temperature=self.temperature,
+            max_tokens=1024
+        )
+        content = response.choices[0].message.content.strip()
+        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=content))])
+
+    def _get_message_text(self, messages):
+        if isinstance(messages, list):
+            return " ".join([msg.content for msg in messages])
+        return messages.content
+
+    @property
+    def _llm_type(self):
+        return "chat-groq"
+
+# ✅ Function to return a QA chain
 def create_qa_chain_from_pdf(pdf_path):
     loader = PyPDFLoader(pdf_path)
     documents = loader.load()
     
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     texts = splitter.split_documents(documents)
-    
+
     embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-m3")
     vectorstore = FAISS.from_documents(texts, embeddings)
 
-    llm = ChatGroq(
-        model="llama3-8b-8192",
-        temperature=0.3,
-        api_key=groq_api_key,
-    )
+    llm = ChatGroq(model="llama3-8b-8192", temperature=0.3, api_key=groq_api_key)
 
     qa_chain = RetrievalQA.from_chain_type(
         llm=llm,