langchat / chain.py
Tanmay09516's picture
Update chain.py
242ea7d verified
# chain.py
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
def init_conversational_chain(llm, retriever):
"""
Initialize the Conversational Retrieval Chain with memory and custom prompt.
Args:
llm: The language model to use.
retriever: The retriever to fetch relevant documents.
Returns:
An instance of ConversationalRetrievalChain.
"""
# Initialize conversation memory
memory = ConversationBufferMemory(
return_messages=True,
memory_key="chat_history",
output_key="answer"
)
# Define a custom prompt template
custom_prompt = PromptTemplate(
input_variables=["context", "question"],
template=(
"You are LangChat, a knowledgeable assistant for the LangChain Python Library. "
"Given the following context from the documentation, provide a helpful answer to the user's question. \n\n"
"Context:\n{context}\n\n"
"You can ignore the context if the question is a simple chat like Hi, hello, and just respond in a normal manner as LangChat, otherwise use the context to answer the query."
"If you can't find the answer from the sources, mention that clearly instead of making up an answer.\n\n"
"Question: {question}\n\n"
"Answer:"
)
)
# Initialize the Conversational Retrieval Chain
qa_chain = ConversationalRetrievalChain.from_llm(
llm=llm,
retriever=retriever,
memory=memory,
return_source_documents=True,
combine_docs_chain_kwargs={"prompt": custom_prompt},
verbose=False
)
return qa_chain