Update app.py

app.py

@@ -7,7 +7,6 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from sentence_transformers import SentenceTransformer
-from langchain import hub
 import bs4
 import torch
 
@@ -114,6 +113,11 @@ class CustomLanguageModel:
         # Implement logic to generate a response based on prompt and context
         return f"Generated response based on prompt: '{prompt}' and context: '{context}'."
 
+# Define a callable class for RAGPrompt
+class RAGPrompt:
+    def __call__(self, data):
+        return {"question": data["question"], "context": data["context"]}
+
 # Submit button for chat
 if st.button("Submit Query"):
     if not query:
@@ -141,23 +145,27 @@ if st.button("Submit Query"):
 
         # Retrieve and generate using the relevant snippets of the blog
         retriever = vectorstore.as_retriever()
-        prompt = hub.pull("rlm/rag-prompt")
 
+        # Retrieve relevant documents
+        retrieved_docs = retriever.get_relevant_documents(query)
+
+        # Format the retrieved documents
         def format_docs(docs):
             return "\n\n".join(doc.page_content for doc in docs)
 
+        context = format_docs(retrieved_docs)
+
         # Initialize the language model
         custom_llm = CustomLanguageModel()
 
-        # Extract context from the retriever
-        retrieved_docs = retriever.get_relevant_documents(query)
-        context = format_docs(retrieved_docs)
+        # Initialize RAG chain using the prompt
+        prompt = RAGPrompt()
 
         rag_chain = (
-            {"context": context, "question": query}
-            | prompt
-            | (lambda data: custom_llm.generate(data["question"], data["context"]))
-            | StrOutputParser()
+            {"context": context, "question": query}  # Start the chain with context and question
+            | prompt                                 # Use the custom prompt
+            | (lambda data: custom_llm.generate(data["question"], data["context"]))  # Pass question and context to LLM
+            | StrOutputParser()  # Parse the output
         )
 
         # Generate the answer using the user's query