RAG项目笔记

• 说明：打算先用 langchain 库做一个简单的rag系统，然后用成熟的框架搭建一个可用于生产环境的项目。

模型

• 模型的获取此处用 langchain_openai 库，通过 url ， api ， model 来达成通用，而不是去找模型厂商提供的第三方包，当然，直接用 openai 库来获取模型也可以，与 langchain 是可以协作的。

# config.py
# model
import os

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
load_dotenv()

chat_api_key = os.getenv("CHAT_API_KEY") # 这里需要建立.env文件，在里面填入key，或者这里直接替换也可以
chat_model_name = "deepseek-chat"
chat_model_url = "https://api.deepseek.com/v1"

embedding_api_key = os.getenv("EMBEDDING_API_KEY") 
embedding_model_name = "embedding-3"
embedding_model_url = "https://open.bigmodel.cn/api/paas/v4"

chat_model = ChatOpenAI(
    model = chat_model_name,
    base_url = chat_model_url,
    api_key = chat_api_key
)

embedding_model = OpenAIEmbeddings(
    model = embedding_model_name,
    base_url = embedding_model_url,
    api_key = embedding_api_key
)

• 自己调整参数就完事了，不想用env的也可以直接换掉，不影响，只要key有值就行。
• 但是所选的模型需要兼容openai格式才行，不然就只能自己改了。

文件上传服务

• 然后是文件上传服务，这里使用 streamlit 库实现前端页面与上传服务。
• 使用教程请查看streamlit中文文档，此处不再多说：(https://www.aidoczh.com/streamlit)。
• 读者搞个其它实现也可以，我的代码其实美化居多，不管什么实现，核心功能只需要获取上传的文件对象即可。

# file_uploader.py
import streamlit as st
import time
from knowledge_base import KnowledgeBase

# 1.设置页面属性（也就是页面标签显示的图标与文字）
# 2.页面标题
# 3.初始化页面永久属性：文件计数，引用对象等
# 4.界面美化（可选）
# 5.完成文件上传功能
# 6.完成文件处理功能（可选，主要是记录文件并显示出来）

st.set_page_config(
    page_title="RAG文件上传系统",
    page_icon="📚"
)

st.title("📚RAG文件上传系统")
st.markdown("- 上传文本文件，为RAG系统准备数据")

# 记录已上传文件信息
if "uploaded_files" not in st.session_state.keys():
    st.session_state["uploaded_files"] = []

# 记录已上传文件名，方便查重
if "files_name" not in st.session_state.keys():
    st.session_state["files_name"] = []

# 记录唯一标识符，刷新st.file_uploader状态
if "key" not in st.session_state.keys():
    st.session_state["key"] = 0

# 记录KnowledgeBase对象
if "KnowledgeBase" not in st.session_state.keys():
    st.session_state["KnowledgeBase"] = KnowledgeBase()


with st.sidebar:
    st.header("设置")

    file_type = st.multiselect(
        label="请选择文件类型",
        options=["txt", "md", "csv", "json"],
        default=["txt", "json"]
    )

    max_size = st.slider(
        label="请选择文件最大大小（MB）",
        min_value=1, max_value=50, value=10
    )

    if st.button(label="页面刷新"):
        st.rerun()

    st.divider()
    
    st.header("使用说明")
    st.write("1. 选择文件类型和大小")
    st.write("2. 拖拽或选择文件上传")
    st.write("3. 查看上传的文件列表")
    st.write("4. 点击处理文件按钮继续")

st.header("上传文件")

# 通过key来控制刷新，但是会有延迟bug
uploaded_files = st.file_uploader(
    "请选择文件（目前有个小的延迟bug，请上传之后点击刷新按钮刷新页面）",
    type=file_type,
    accept_multiple_files=True,
    key=st.session_state["key"],
    help="选择完成后点击'确认上传'"
)

if uploaded_files:
    with st.spinner("正在上传..."):
        for uploaded_file in uploaded_files:
            try:
                uploaded_file_name = uploaded_file.name
                uploaded_file_size = uploaded_file.size / 1024
                uploaded_file_type = uploaded_file.type
                
                if uploaded_file_name in st.session_state["files_name"]:
                    st.warning(f"文件 {uploaded_file_name} 已存在，跳过处理")
                    continue
                if uploaded_file_size > max_size * 1024:
                    st.warning(f"文件 {uploaded_file_name} 过大，跳过处理")
                    continue
                
                uploaded_file_content = uploaded_file.read().decode("utf-8").strip()
                uploaded_file_info = {
                    'name': uploaded_file_name,
                    'size': uploaded_file_size + 'KB',
                    "type": uploaded_file_type,
                    'content': uploaded_file_content,
                    'upload_time': time.strftime("%Y-%m-%d %H:%M:%S")
                }

                st.session_state["uploaded_files"].append(uploaded_file_info)
                st.session_state["files_name"].append(uploaded_file_name)
                
                # 处理添加到知识库服务
                print(st.session_state["KnowledgeBase"].upload_file(uploaded_file_info))
                
            except Exception as e:
                st.error(f"文件 {uploaded_file.name} 读取失败: {str(e)}")
        
        st.session_state["key"] += 1
        time.sleep(1)

if st.session_state["uploaded_files"]:
    st.subheader("已上传文件：")
    for file_info in st.session_state["uploaded_files"]:
        st.write(f"文件名：{file_info['name']}")
        st.write(f"文件大小：{file_info['size']:.2f}KB")
        st.write(f"文件类型：{file_info['type']}")
        st.write(f"文件内容：{file_info['content'][:300] + '...' if len(file_info['content']) > 300 else file_info['content']}")
        st.write(f"文件上传时间：{file_info['upload_time']}")
        st.divider()

知识库服务

• 需要实现把上传的文件转化为向量存储的功能。
• 首先要查重，查重通过md5实现，维护一个存储md5值的文件即可。
• 然后要把大文件切分，通过 langchain_text_splitters 自带功能实现。
• 然后向量化存储，通过 langchain_chroma 实现。
• 还需要提供查询服务，同样通过 langchain_chroma 自带方法可实现。

# config.py
# text_splitter
from langchain_text_splitters import RecursiveCharacterTextSplitter

separators_ = ["\n\n", "\n", ".", "。", "!", "！", "?", "？", " ", ""] # 分割符
chunk_size_ = 1000 # 文本块大小
chunk_overlap_ = 100 # 允许文本块重叠区域大小
length_function_ = len

text_splitter = RecursiveCharacterTextSplitter(
    separators=separators_,
    chunk_size=chunk_size_,
    chunk_overlap=chunk_overlap_,
    length_function=length_function_,
)



# chroma
from langchain_chroma import Chroma

collection_name_ = "RAG_DATABASE"
embedding_function_ = embedding_model
persist_direction_ = pathlib.Path(__file__).parent / "chroma_bd"
md5_path = pathlib.Path(__file__).parent / "md5.txt" # 自己建立
top_k = 2 # 查找文档时选择top_k个最相近文档

os.makedirs(name=persist_direction_, exist_ok=True)

chroma = Chroma(
    collection_name=collection_name_,
    embedding_function=embedding_function_,
    persist_directory=persist_direction_,
)

retriever = chroma.as_retriever(search_kwargs={'k': top_k})

# knowledge_base.py
# 实现知识库类
# 主要功能：把上传的文件转化为向量存储起来
# 附带功能：通过MD5查重

import config
import hashlib

class KnowledgeBase():
    def __init__(self):
        self.text_splitter = config.text_splitter
        self.chroma = config.chroma

    def get_knowledge(self, text):
        retriever = config.retriever
        return retriever.invoke(text)

    def upload_file(self, file):
        if not isinstance(file, dict):
            raise Exception("KnowledgeBase的upload_file方法需要字典对象！")
        
        # 读取内容
        content = file['content']
        meta_data = {
            'name': file['name'],
            'size': file['size'],
            "type": file['type'],
            'upload_time': file['upload_time']
        }
        # 查重
        if self.check_md5(content):
            return f"文件 {file['name']} 已存在"
        
        # 切割
        if len(content) > config.chunk_size_:
            chunks = self.text_splitter.split_text(content)
        else:
            chunks = [content]

        # 向量化存储
        self.chroma.add_texts(
            texts=chunks,
            metadatas=[meta_data for _ in chunks]
        )

        # 存入md5
        self.save_md5(content)

        return f"文件 {file['name']} 已存入向量库"
    
    def check_md5(self, content):
        if not isinstance(content, str):
            raise Exception("KnowledgeBase的check_md5方法需要字符串对象！")
        with open(file=config.md5_path, mode='r', encoding='utf-8') as f:
            md5_str = self.get_md5(content)
            for line in f.readlines():
                if md5_str == line.strip():
                    return True
            return False

    def save_md5(self, content):
        if not isinstance(content, str):
            raise Exception("KnowledgeBase的save_md5方法需要字符串对象！")
        with open(file=config.md5_path, mode='a', encoding='utf-8') as f:
            md5_str = self.get_md5(content)
            f.write(md5_str + '\n')

    def get_md5(self, content):
        if not isinstance(content, str):
            raise Exception("KnowledgeBase的get_md5方法需要字符串对象！")
        content = content.encode(encoding='utf-8')
        md5_obj = hashlib.md5()
        md5_obj.update(content)
        return md5_obj.hexdigest()

历史消息存储

• 历史消息存储可以通过把每次的聊天记录都存储为文件来实现，每次对话时带上历史记录就行。
• 实现一个继承 BaseChatMessageHistory 的类（其实继不承继承无所谓，压根不影响功能，只要实现了相应方法就行）。
• prompt模板使用 MessagePlaceHolder 作为历史消息占位符，这样可以保证对话结构不被破坏。
• 我的实现没有用到太多的chainlang组件，因为这样比较好扩展（RAG部分也是）。

# history_store.py
import json
from typing import Sequence
import os
from langchain_core.messages import BaseMessage, messages_to_dict, messages_from_dict

class HistoryStore():
    def __init__(self, user_id, chat_id, storage_path):
        self.user_id = user_id
        self.chat_id = chat_id
        self.storage_path = storage_path

        self.file_path = os.path.join(storage_path, f"{user_id}")

        self.file = os.path.join(self.file_path, f"{chat_id}.json")
        os.makedirs(self.file_path, exist_ok=True)
    
    def add_message(self, messages: Sequence[BaseMessage]) -> None:
        if not self.messages:
            all_messages = messages
        else:
            all_messages = list(self.messages)
            all_messages.extend(messages)

        new_messages = messages_to_dict(all_messages)

        with open(self.file, 'w', encoding='utf-8') as f:
            json.dump(new_messages, f)

    @property
    def messages(self) -> list[BaseMessage]:
        try:
            with open(self.file, 'r', encoding='utf-8') as f:
                if f is not None:
                    return messages_from_dict(json.load(f))
                else:
                    return []
        except Exception:
            return []
    
    def clear(self) -> None:
        with open(self.file, 'w', encoding='utf-8') as f:
            json.dump([], f)

RAG

• RAG的理念就是在问问题的时候附带上与问题相关的信息（虽然最近似乎有研究表明ai压根不会读上下文）。
• 首先需要提示词模板来做到提示词嵌入。
• 然后检索文档，之后和问题一起放入提示词模板。
• 最后给ai，并解析ai的回答。

# config.py
# RAG
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

prompt_template = ChatPromptTemplate.from_messages(
    [
        ('system', '以已知的参考资料为主，简洁并专业地回答用户问题。参考资料：\n{context}'),
        MessagesPlaceholder('history'),
        ('user', '请回答用户问题：{input}')
    ]
)



# history_store
storage_path = pathlib.Path(__file__).parent / "history_store"

# rag.py
import config
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import AIMessage, HumanMessage
from history_store import HistoryStore

class RAG():
    def __init__(self, user_id=1, chat_id=1):
        # 提示词模板（要修改到config里改）
        self.prompt_template = config.prompt_template

        # 知识库查询
        self.retriever = config.retriever

        # 历史记录
        self.history = HistoryStore(user_id, chat_id, config.storage_path)

        # 模型
        self.chat_model = config.chat_model

    def get_chain(self, input):
        dic = {
            'input': input,
            'context': self.format_doc(self.retriever.invoke(input)),
            'history': self.history.messages,
        }
        # ChatOpenAI会自动处理传入的提示词模板类型，所以不需要转化为字符串
        chain = (
            self.prompt_template | self.show_str | self.chat_model | StrOutputParser()
        )

        result = chain.invoke(dic)

        message = [
            HumanMessage(input), 
            AIMessage(result)
        ]

        self.history.add_message(message)

        return result

    def show_str(self, text):
        print('=='*20)
        print(text.to_string())
        print('=='*20)
        return text

    def format_doc(self, docs):
        if not isinstance(docs, list):
            raise Exception("RAG的format_doc函数需要list类型参数！")
        # 空列表是false
        if not docs:
            return "无相关资料"
        format_str = ""
        for doc in docs:
            format_str += f"文档元数据：{doc.metadata}\n文档内容：{doc.page_content}\n\n"
        return format_str

问答界面

• 最后依旧用streamlit做一个简单的问答界面即可

# chat_interface.py
import streamlit as st
from rag import RAG

st.set_page_config(
    page_title="RAG问答系统",
    page_icon="📚"
)

st.title("📚RAG问答系统")
st.divider()

if 'messages' not in st.session_state.keys():
    st.session_state['messages'] = [{'name':'ai', 'content':'请问有什么可以帮助你的？'}]

if 'RAG' not in st.session_state.keys():
    st.session_state['RAG'] = RAG(user_id=1, chat_id=1)

for message in st.session_state['messages']:
    st.chat_message(message['name']).write(message['content'])

prompt = st.chat_input()

if prompt:
    st.chat_message('user').write(prompt)
    st.session_state['messages'].append({'name':'user', 'content':prompt})

    with st.spinner("AI思考中..."):
        result = st.session_state['RAG'].get_chain(prompt)
        st.chat_message('ai').write(result)
        st.session_state['messages'].append({'name':'ai', 'content':result})

一些函数介绍

• MessagesPlaceholder ：可以自动解析传入的messages格式，而不是直接变成字符串，比如对话历史，用这个会保留对话历史的格式，如果直接用占位符，那就是把聊天历史直接转化为字符串全拼一起了
• to_string : prompt出来的数据类型是message，直接print不方便看，这时候直接 to_string 就好了。