基于langchain的简单RAG的实现

news2025/6/9 0:37:16

闲来无事，想研究一下RAG的实现流程，看网上用langchain的比较多，我自己在下面也跑了跑，代码很简单，以次博客记录一下，方便回顾

langchain

LangChain 是一个基于大型语言模型（LLM）开发应用程序的框架。LangChain 简化了LLM应用程序生命周期的每个阶段。

比如，在下面的实现中，LangChain可以将LLM、提示词模板、检索器组合在一起快速的完成检索增强整个流程，而不需要你去关心底层具体是怎么实现的。

代码demo

实现思路：

加载文档，并对文档进行切分
将切分后的文档转化为向量，存储到向量库中
根据用户query去向量库中检索，找到最相关的回复，并拼接到prompt中
根据最新的prompt调用大模型产生增强回复

加载文档 -> 切分文档 -> 创建向量数据库 -> 执行相似度搜索 -> 构建并增强 prompt -> 使用模型生成回答

import os
from openai import OpenAI
import requests
from langchain.text_splitter import CharacterTextSplitter
from weaviate.embedded import EmbeddedOptions
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser

from langchain_community.document_loaders import TextLoader
from langchain_community.chat_models import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document
from langchain_core.messages import HumanMessage, SystemMessage


"""
实现一
"""
def method_one(vectorstore,llm,query):

    # 根据用户query进行检索，并将检索结果拼接到prompt中
    def augment_prompt(query: str):

        # 获取top2的文本片段
        results = vectorstore.similarity_search(query, k=1)
        source_knowledge = "\n\n".join([x.page_content for x in results])

        # 构建prompt
        augmented_prompt = f"""你叫david,你需要解答xxx问题

        ###参考样例###
        {source_knowledge}
        """
        return augmented_prompt

    prompt=augment_prompt(query)
    print(prompt)

    # 封装输入
    messages = [
        SystemMessage(content=prompt),
        HumanMessage(content=query),
    ]

    # 生成检索增强回复
    res = llm.invoke(messages)

    return res.content


"""
实现二
"""
def method_two(vectorstore,llm,query):
    
    # 将 vectorstore 转换为一个检索器
    retriever = vectorstore.as_retriever()

    # 定义提示模板
    template = """你叫david,你需要解答xxx问题

    ###参考样例###
    {context}

    ###用户问题###
    {question}
    """

    prompt = ChatPromptTemplate.from_template(template)

    print(prompt)

    # LangChain 提供了一个高度模块化和可组合的框架就是链，使得你可以根据任务的特性自定义每个组件，并将它们按需组合成执行流程
    # 定义一个执行流程链，包含如下组件
    # {"context": retriever,  "question": RunnablePassthrough()}：用来将上下文（通过检索器获得）和用户问题传递给后续组件
    # prompt里面的占位符与上述定义的context和question是要保持一致的
    # StrOutputParser()：该组件用于解析模型的输出，将其转换为字符串格式
    rag_chain = (
        {"context": retriever,  "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    response = rag_chain.invoke(query)

    return response



if __name__=="__main__":

    # 加载单个文档，这里只需要匹配单个文档里面的片段
    path="../rag/faq.txt"
    loader = TextLoader(path)
    documents = loader.load()

    # 如果需要加载多个文档，将上述path改为跟路径即可，然后通过下述两行代码对多个文档进行切分
    # text_splitter = CharacterTextSplitter()
    # doc = text_splitter.split_documents(documents)

    # 切分文档，给定的文档内容主要是通过换行符分隔的
    text = documents[0].page_content
    chunks = [Document(page_content=chunk) for chunk in text.split("\n\n\n") if chunk.strip()]

    #  将文档片段转化为向量，并存储到 
    # Chroma 是一个 开源的向量数据库，用于存储和检索向量嵌入
    model_name = "../model/bge-base-zh-v1.5"
    embedding = HuggingFaceEmbeddings(model_name=model_name)
    vectorstore_hf = Chroma.from_documents(documents=chunks, embedding=embedding , collection_name="huggingface_embed")
    vectorstore = Chroma.from_documents(chunks, embedding)

    # 初始化对话模型
    llm = ChatOpenAI(
        openai_api_key="",
        openai_api_base="",
        model='qwen-max'
    )
    
    # 用户query
    query = "今天天气如何？"

    # 检索增强之后的回答
    enhanced_result=method_one(vectorstore,llm,query)
    # enhanced_result=method_two(vectorstore,llm,query)

    print(enhanced_result)