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

# 定义一个流输出解析器类，用于模拟流式输出效果
class StreamOutputParser(StrOutputParser):
    def parse(self, text: str):
        chunk_size = 1
        # 逐字输出文本内容，模拟流式输出的效果
        for i in range(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")

# 格式化文档内容的函数
def format_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 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
   

   这些库提供了网页加载、文档分割、向量存储和语言模型调用的工具。

2. 设置 OpenAI API 密钥

   os.environ["OPENAI_API_KEY"] = getpass.getpass()
   

   通过getpass.getpass()获取用户输入的API密钥，以避免在代码中直接暴露密钥。

3. 定义流输出解析器类

   class StreamOutputParser(StrOutputParser):
       def parse(self, text: str):
           chunk_size = 1
           for i in range(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链

   def format_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链检索并生成关于“任务分解”的回答，并以流式输出的方式逐字打印结果。