Update app.py

app.py

@@ -111,9 +111,8 @@ if 'chat_history' not in st.session_state:
 # CustomLanguageModel class with proper context argument
 class CustomLanguageModel:
     def generate(self, prompt, context):
-        # This function should handle both prompt and context
-        # For now, we return a placeholder response for demo purposes
-        return f"Generated response based on '{prompt}' and context provided."
+        # Implement logic to generate a response based on prompt and context
+        return f"Generated response based on prompt: '{prompt}' and context: '{context}'."
 
 # Submit button for chat
 if st.button("Submit Query"):
@@ -150,15 +149,19 @@ if st.button("Submit Query"):
         # Initialize the language model
         custom_llm = CustomLanguageModel()
 
+        # Extract context from the retriever
+        retrieved_docs = retriever.retrieve(query)
+        context = format_docs(retrieved_docs)
+
         rag_chain = (
-            {"context": retriever | format_docs, "question": RunnablePassthrough()}
+            {"context": context, "question": query}
             | prompt
-            | custom_llm.generate  # Adjust based on actual usage
+            | (lambda data: custom_llm.generate(data["question"], data["context"]))
             | StrOutputParser()
         )
 
         # Generate the answer using the user's query
-        result = rag_chain.invoke(query)
+        result = rag_chain.invoke({"question": query, "context": context})
 
         # Store query and response in session for chat history
         st.session_state['chat_history'].append((query, result))