LangChain实现文档检索和增强生成的示例

在这篇博客中，我们将介绍如何使用LangChain和Chroma来实现文档的检索和增强生成。我们将以一个具体的实例来展示这一过程，具体代码如下（修改自官方文档实例）：

import getpass
import os
import time

import bs4
from langchain import hub
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 设置 OpenAI API 密钥# 使用 getpass.getpass() 获取用户输入的 API 密钥，避免直接在代码中暴露密钥
os.environ["OPENAI_API_KEY"]= getpass.getpass()from langchain_openai import ChatOpenAI

# 定义一个流输出解析器类，用于模拟流式输出效果classStreamOutputParser(StrOutputParser):defparse(self, text:str):
        chunk_size =1# 逐字输出文本内容，模拟流式输出的效果for i inrange(0,len(text), chunk_size):yield text[i:i + chunk_size]
            time.sleep(0.1)# 模拟流式输出的延迟# 初始化语言模型
llm = ChatOpenAI(model="gpt-3.5-turbo-0125")# 加载、分块和索引博客内容
loader = WebBaseLoader(
    web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(
            class_=("post-content","post-title","post-header"))),)# 加载指定网页的内容
docs = loader.load()# 使用递归字符分割器对文档进行分块# chunk_size 指定每块的最大字符数，chunk_overlap 指定相邻块之间的重叠字符数
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)# 创建一个 Chroma 向量存储来存储文档
vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())# 使用向量存储检索并生成相关的博客片段
retriever = vectorstore.as_retriever()# 从Hub中获取一个RAG（检索增强生成）提示模板
prompt = hub.pull("rlm/rag-prompt")# 格式化文档内容的函数defformat_docs(docs):return"\n\n".join(doc.page_content for doc in docs)# 构建 RAG（检索增强生成）链
rag_chain =({"context": retriever | format_docs,"question": RunnablePassthrough()}| prompt
    | llm
    | StreamOutputParser())# 检索并以流式输出方式打印结果
result_stream = rag_chain.invoke("What is Task Decomposition?使用中文回答。")for chunk in result_stream:print(chunk, end='', flush=True)

项目依赖

pip install flask bs4 langchain langchain_chroma langchain_community langchain_core langchain_openai

启动

需科学上网，Password输入chatgpt秘钥

代码剖析

1. 导入必要的库import getpassimport osimport timeimport bs4from langchain import hubfrom langchain_chroma import Chromafrom langchain_community.document_loaders import WebBaseLoaderfrom langchain_core.output_parsers import StrOutputParserfrom langchain_core.runnables import RunnablePassthroughfrom langchain_openai import OpenAIEmbeddingsfrom langchain_text_splitters import RecursiveCharacterTextSplitter这些库提供了网页加载、文档分割、向量存储和语言模型调用的工具。
2. 设置 OpenAI API 密钥os.environ["OPENAI_API_KEY"]= getpass.getpass()通过getpass.getpass()获取用户输入的API密钥，以避免在代码中直接暴露密钥。
3. 定义流输出解析器类classStreamOutputParser(StrOutputParser):defparse(self, text:str): chunk_size =1for i inrange(0,len(text), chunk_size):yield text[i:i + chunk_size] time.sleep(0.1)这个类逐字输出文本内容，并模拟流式输出的延迟效果。
4. 初始化语言模型llm = ChatOpenAI(model="gpt-3.5-turbo-0125")初始化一个使用GPT-3.5-turbo模型的语言模型实例，用于后续的文本生成。
5. 加载和分块文档loader = WebBaseLoader( web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",), bs_kwargs=dict( parse_only=bs4.SoupStrainer( class_=("post-content","post-title","post-header"))),)docs = loader.load()text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)splits = text_splitter.split_documents(docs)使用WebBaseLoader加载指定网页的内容，并通过RecursiveCharacterTextSplitter将文档分割成较小的块。
6. 创建向量存储和检索器vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())retriever = vectorstore.as_retriever()prompt = hub.pull("rlm/rag-prompt")通过Chroma创建一个向量存储，将分割后的文档存储起来，并设置一个检索器用于检索相关文档内容。同时，从Hub中获取一个RAG（检索增强生成）提示模板。
7. 构建RAG链defformat_docs(docs):return"\n\n".join(doc.page_content for doc in docs)rag_chain =({"context": retriever | format_docs,"question": RunnablePassthrough()}| prompt | llm | StreamOutputParser())定义一个函数format_docs来格式化文档内容，并构建一个RAG链，以实现从检索到生成的完整流程。
8. 检索并以流式输出方式打印结果result_stream = rag_chain.invoke("What is Task Decomposition?使用中文回答。")for chunk in result_stream:print(chunk, end='', flush=True)通过RAG链检索并生成关于“任务分解”的回答，并以流式输出的方式逐字打印结果。