File size: 721 Bytes
362de84 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
"""Conversational QA chain setup."""
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
def create_qa_chain(llm, retriever):
"""
Create conversational QA chain with memory.
Args:
llm: Language model
retriever: Document retriever
Returns:
ConversationalRetrievalChain
"""
memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True,
output_key="answer"
)
chain = ConversationalRetrievalChain.from_llm(
llm=llm,
retriever=retriever,
memory=memory,
return_source_documents=True
)
return chain
|