Update rag_chain.py

rag_chain.py

@@ -1,15 +1,41 @@
-from langchain_community.chains import RetrievalQA
 from langchain_groq import ChatGroq
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 
 def build_rag_chain(vectorstore, groq_api_key):
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
+
     llm = ChatGroq(
         api_key=groq_api_key,
         model="llama3-8b-8192",
         temperature=0
     )
 
-    return RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
-        return_source_documents=True
+    prompt = ChatPromptTemplate.from_template(
+        """
+        You are an expert software engineer.
+        Answer the question using ONLY the context below.
+
+        Context:
+        {context}
+
+        Question:
+        {question}
+        """
     )
+
+    def format_docs(docs):
+        return "\n\n".join(d.page_content for d in docs)
+
+    chain = (
+        {
+            "context": retriever | format_docs,
+            "question": RunnablePassthrough()
+        }
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+
+    return chain