构建RAG应用
- 1.加载向量数据库
- 2.测试向量库检索答案
- 3.langchain中的gpt做为llm
- 4.检索问答框架【向量库+llm】
- 5.保留回话历史
- 6.根据会话记忆进行回复
- 7.汇总
1.加载向量数据库
如何将数据保存为本地向量库参考上一篇: 基于LangChain框架搭建知识库
def load_vector_db(self):
"""加载向量数据库"""
from langchain.vectorstores.chroma import Chroma
from langchain_openai import OpenAIEmbeddings
embedding = OpenAIEmbeddings(model="text-embedding-3-small",
openai_api_key=self.api_key,
http_client=self.http_client)
# 加载数据库
vectordb = Chroma(
persist_directory=self.persist_directory, # 允许我们将persist_directory目录保存到磁盘上
embedding_function=embedding
)
print(f"向量库中存储的数量：{vectordb._collection.count()}")
return vectordb
2.测试向量库检索答案
def search_txt(self):
"""测试向量库检索问答"""
vectordb = self.load_vector_db()
question = "什么是机器学习"
docs = vectordb.similarity_search(question, k=2)
print(f"检索到的内容数：{len(docs)}")
for i, doc in enumerate(docs):
print(f"检索到的第{i}个内容: \n {doc.page_content}",
end="\n-----------------------------------------------------\n")
3.langchain中的gpt做为llm
def langchain_gpt(self):
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
# data = llm.invoke("请你自我介绍一下自己！")
# print(data)
return llm
4.检索问答框架【向量库+llm】
def llm_response(self):
"""根据向量知识库+llm 检索问答"""
from langchain.prompts import PromptTemplate
llm = self.langchain_gpt()
vectordb = self.load_vector_db()
template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答
案。最多使用三句话。尽量使答案简明扼要。总是在回答的最后说“谢谢你的提问！”。
{context}
问题: {question}
"""
QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"], template=template)
# 基于模版的检索链
from langchain.chains import RetrievalQA
qa_chain = RetrievalQA.from_chain_type(llm,
retriever=vectordb.as_retriever(),
return_source_documents=True,
chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})
question_1 = "什么是南瓜书？"
question_2 = "王阳明是谁？"
result = qa_chain({"query": question_1})
print("大模型+知识库后回答 question_1 的结果：")
print(result["result"])
# 南瓜书是一本由Datawhale编委会编写的书籍，主要内容尚未明确，但可以通过在线阅读地址和PDF获取地址进行阅读。它采用知识共享署名-非商业性使用-相同方式共享4.0 国际许可协议进行许可。谢谢你的提问！
result = qa_chain({"query": question_2})
print("大模型+知识库后回答 question_2 的结果：")
print(result["result"])
# 我不知道王阳明是谁。谢谢你的提问！
prompt_template = """请回答下列问题:
{}""".format(question_1)
### 基于大模型的问答
res = llm.predict(prompt_template)
print(f'llm回答question1：{res}')
prompt_template = """请回答下列问题:
{}""".format(question_2)
### 基于大模型的问答
res = llm.predict(prompt_template)
print(f'llm回答question2：{res}')
5.保留回话历史
def memory_chain(self):
"""保留会话记忆"""
from langchain.memory import ConversationBufferMemory
memory = ConversationBufferMemory(
memory_key="chat_history", # 与 prompt 的输入变量保持一致。
return_messages=True # 将以消息列表的形式返回聊天记录，而不是单个字符串
)
return memory
6.根据会话记忆进行回复
def history_chain(self):
"""根据会话记忆进行问答"""
from langchain.chains import ConversationalRetrievalChain
vectordb = self.load_vector_db()
llm = self.langchain_gpt()
memory = self.memory_chain()
retriever = vectordb.as_retriever()
qa = ConversationalRetrievalChain.from_llm(
llm,
retriever=retriever,
memory=memory
)
# 第一次提问后注释掉
# question = "什么是南瓜书？"
# result = qa({"question": question})
# print(result['answer'])
# 第二次提问，根据上次提问内容进行继续回答
question = "我应该怎么学习南瓜书"
result = qa({"question": question})
print(result['answer'])
7.汇总
import os
from dotenv import load_dotenv, find_dotenv
from openai import OpenAI
import httpx
class RAG:
def __init__(self):
_ = load_dotenv(find_dotenv())
# 获取openai key
self.api_key = os.environ.get("OPENAI_API_KEY")
# 本地向量数据库持久化路径
self.persist_directory = './vector_db/chroma'
# 设置代理地址
self.proxy = 'http://127.0.0.1:7890'
os.environ["HTTPS_PROXY"] = self.proxy
os.environ["HTTP_PROXY"] = self.proxy
def load_vector_db(self):
"""加载向量数据库"""
from langchain.vectorstores.chroma import Chroma
from langchain_openai import OpenAIEmbeddings
embedding = OpenAIEmbeddings(model="text-embedding-3-small",
openai_api_key=self.api_key,
http_client=self.http_client)
# 加载数据库
vectordb = Chroma(
persist_directory=self.persist_directory, # 允许我们将persist_directory目录保存到磁盘上
embedding_function=embedding
)
print(f"向量库中存储的数量：{vectordb._collection.count()}")
return vectordb
def search_txt(self):
"""测试向量库检索问答"""
vectordb = self.load_vector_db()
question = "什么是机器学习"
docs = vectordb.similarity_search(question, k=2)
print(f"检索到的内容数：{len(docs)}")
for i, doc in enumerate(docs):
print(f"检索到的第{i}个内容: \n {doc.page_content}",
end="\n-----------------------------------------------------\n")
def langchain_gpt(self):
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
# data = llm.invoke("请你自我介绍一下自己！")
# print(data)
return llm
def llm_response(self):
"""根据向量知识库+llm 检索问答"""
from langchain.prompts import PromptTemplate
llm = self.langchain_gpt()
vectordb = self.load_vector_db()
template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答
案。最多使用三句话。尽量使答案简明扼要。总是在回答的最后说“谢谢你的提问！”。
{context}
问题: {question}
"""
QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"], template=template)
# 基于模版的检索链
from langchain.chains import RetrievalQA
qa_chain = RetrievalQA.from_chain_type(llm,
retriever=vectordb.as_retriever(),
return_source_documents=True,
chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})
question_1 = "什么是南瓜书？"
question_2 = "王阳明是谁？"
result = qa_chain({"query": question_1})
print("大模型+知识库后回答 question_1 的结果：")
print(result["result"])
# 南瓜书是一本由Datawhale编委会编写的书籍，主要内容尚未明确，但可以通过在线阅读地址和PDF获取地址进行阅读。它采用知识共享署名-非商业性使用-相同方式共享4.0 国际许可协议进行许可。谢谢你的提问！
result = qa_chain({"query": question_2})
print("大模型+知识库后回答 question_2 的结果：")
print(result["result"])
# 我不知道王阳明是谁。谢谢你的提问！
prompt_template = """请回答下列问题:
{}""".format(question_1)
### 基于大模型的问答
res = llm.predict(prompt_template)
print(f'llm回答question1：{res}')
prompt_template = """请回答下列问题:
{}""".format(question_2)
### 基于大模型的问答
res = llm.predict(prompt_template)
print(f'llm回答question2：{res}')
def memory_chain(self):
"""保留会话记忆"""
from langchain.memory import ConversationBufferMemory
memory = ConversationBufferMemory(
memory_key="chat_history", # 与 prompt 的输入变量保持一致。
return_messages=True # 将以消息列表的形式返回聊天记录，而不是单个字符串
)
return memory
def history_chain(self):
"""根据会话记忆进行问答"""
from langchain.chains import ConversationalRetrievalChain
vectordb = self.load_vector_db()
llm = self.langchain_gpt()
memory = self.memory_chain()
retriever = vectordb.as_retriever()
qa = ConversationalRetrievalChain.from_llm(
llm,
retriever=retriever,
memory=memory
)
# 第一次提问后注释掉
# question = "什么是南瓜书？"
# result = qa({"question": question})
# print(result['answer'])
# 第二次提问，根据上次提问内容进行继续回答
question = "我应该怎么学习南瓜书"
result = qa({"question": question})
print(result['answer'])
if __name__ == '__main__':
rag = RAG()
# 基于向量库+llm进行检索问答
# rag.llm_response()
# 基于llm+历史会话记忆进行问答
rag.history_chain()
