解决Chainlit中用户会话链对象的正确存取方法

@cl.on_message
async def main(message):
    chain = cl.user_session.set("chain") # 错误：这里应该使用 get
    # ... 后续代码

chain = cl.user_session.get("chain") # 正确：从会话中获取已存储的链对象

from langchain.prompts import PromptTemplate
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores.faiss import FAISS
from langchain.llms import CTransformers
from langchain.chains import RetrievalQA
import chainlit as cl
import os

# 确保DB_FAISS_PATH存在，这里假设它在当前脚本的同级目录
# 实际应用中，应确保此路径有效且包含预训练的FAISS索引
DB_FAISS_PATH = "vectorstores/db_faiss" 
# 确保llama-2-7b-chat.ggmlv3.q8_0.bin模型文件可访问
# 通常放置在与脚本同级或指定路径下
LLM_MODEL_PATH = "llama-2-7b-chat.ggmlv3.q8_0.bin"


custom_prompt_template = """Use the following pieces of information to answer the user's question.
If you don't know the answer, please just say that you don't know the answer, don't try to make up
an answer.

Context: {context}
Question: {question}

Only returns the helpful answer below and nothing else.
Helpful answer:

"""

def set_custom_prompt():
    """
    为QA检索设置PromptTemplate。
    """
    prompt = PromptTemplate(template=custom_prompt_template, 
                            input_variables=['context', 'question'])
    return prompt


def load_llm():
    """
    加载CTransformers模型。
    """
    print("*** Start loading the LLM.")
    llm = CTransformers(
        model=LLM_MODEL_PATH,
        model_type="llama",
        max_new_tokens=512,
        temperature=0.5
    )
    print("****** Finished loading the LLM.")
    return llm

def retrieval_qa_chain(llm, prompt, db):
    """
    创建并返回RetrievalQA链。
    """
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=db.as_retriever(search_kwargs={'k': 2}),
        return_source_documents=True,
        chain_type_kwargs={'prompt': prompt}
    )
    return qa_chain

def qa_bot():
    """
    初始化并返回一个完整的QA机器人链。
    """
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2',
                                       model_kwargs={'device': 'cpu'})

    # 检查FAISS向量库路径是否存在
    if not os.path.exists(DB_FAISS_PATH):
        raise FileNotFoundError(f"FAISS vector store not found at {DB_FAISS_PATH}. Please ensure it is created.")

    db = FAISS.load_local(DB_FAISS_PATH, embeddings)
    print("*******FAISS.load_local() works well.")

    # 检查LLM模型文件是否存在
    if not os.path.exists(LLM_MODEL_PATH):
        raise FileNotFoundError(f"LLM model not found at {LLM_MODEL_PATH}. Please ensure the model file is present.")

    llm = load_llm()
    print("****** LLM loading step works well.")

    qa_prompt = set_custom_prompt()
    qa = retrieval_qa_chain(llm, qa_prompt, db)
    print("****** QA chain creation step works well.")

    return qa

# chainlit ####
@cl.on_chat_start
async def start():
    """
    聊天开始时初始化机器人并将其存储在用户会话中。
    """
    chain = qa_bot()
    msg = cl.Message(content="Starting the bot......")    
    await msg.send()
    msg.content = "Hi, Welcome to the Medical Bot. What is your query?"
    await msg.update()

    # 将初始化的链对象存储到用户会话中
    cl.user_session.set('chain', chain)

@cl.on_message
async def main(message):
    """
    处理用户消息，从会话中获取链对象并执行查询。
    """
    # 从用户会话中获取之前存储的链对象
    chain = cl.user_session.get("chain")

    if chain is None:
        await cl.Message(content="Error: The bot chain was not initialized correctly. Please restart the chat.").send()
        return

    cb = cl.AsyncLangchainCallbackHandler(
        stream_final_answer=True, answer_prefix_tokens=["FINAL", "ANSWER"]
    )
    cb.answer_reached = True # 确保回调在答案生成后触发

    # 假设message对象有一个content属性，或者可以直接作为查询使用
    res = await chain.acall(message.content, callbacks=[cb]) # 传递message.content作为查询
    answer = res["result"]
    sources = res.get("source_documents", []) # 使用.get()以防没有源文档

    if sources:
        # 将源文档格式化为更易读的字符串
        source_texts = [doc.page_content for doc in sources]
        answer += f"\n\n**Sources:**\n" + "\n".join(source_texts)
    else:
        answer += f"\n\n**No Sources Found**"

    await cl.Message(content=answer).send()