一个简单的RAG知识库

导入相关依赖

• 第一步

import os
import streamlit as st
import tempfile
#  会话存储
from langchain.memory import ConversationBufferMemory
#  聊天历史记录
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
#  TXT加载器
from langchain_community.document_loaders import TextLoader
#  向量数据库
from langchain_openai import OpenAIEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
#  向量数据库
# from chromadb import Client
from langchain_community.vectorstores import FAISS as Chroma
# from langchain_chroma import Chroma 
#  提示词模板from langchain_core.prompts import PromptTemplate
#  文本分块
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.agents import create_react_agent, AgentExecutor
# 数据毁掉
from langchain_community.callbacks.streamlit import StreamlitCallbackHandler
from langchain_core.prompts import PromptTemplate
from pydantic import SecretStr
from langchain.tools.retriever import create_retriever_tool

• 第二步 构建LLM

# 初始化模型
DASHSCOPE_API_KEY = SecretStr(os.getenv("DASHSCOPE_API_KEY", "自己的KEY"))

if not DASHSCOPE_API_KEY.get_secret_value():
    raise ValueError("DASHSCOPE_API_KEY 环境变量未设置")

# 初始化 OpenAI 客户端
llm = ChatOpenAI(
    api_key=DASHSCOPE_API_KEY,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo" 
)

• 第三步 使用 python的streamlit 构建简易页面

# 设置标题和页面布局
st.set_page_config(page_title="CRM常见问题", layout="wide")
st.title("CRM常见问题")
# 可以上传多个TXT文件
uploader_file = st.sidebar.file_uploader("上传文件", type=["txt"], accept_multiple_files=True)

if not uploader_file:
    st.info("请上传文件")
    st.stop()

• 第四部构建会话信息

# 如果聊天记录为空 或者清空了聊天记录【session为空】，则创建一个空的历史记录
if "messages" not in st.session_state or st.sidebar.button("清空聊天记录"):
    st.session_state["messages"] = [{"role":"assistant", "content":"你好！请输入问题"}]


# 记载历史聊天记录
for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

# 创建检索工具
from langchain.tools.retriever import create_retriever_tool

tool = create_retriever_tool(
    retriever,
    name="search_docs",
    description="用于搜索文档，输入为问题，返回答案"
)
tools = [tool]

# 创建历史消息对话对象
msgs = StreamlitChatMessageHistory()

# 创建对话缓存区
memory = ConversationBufferMemory(chat_memory=msgs, return_messages=True, memory_key="chat_history",output_key="output")

• 五 构建检索

# 实现检索题
@st.cache_resource(ttl="1h")
def config_retriever(uploader_file):
    # 文档分快对象
    docs = []
    ##temp_dir = tempfile.TemporaryDirectory()  # 更通用
    # 临时文件地址
    temp_dir = tempfile.TemporaryDirectory(dir="/Users/Desktop/")
    for file in uploader_file:
        temp_file_path = os.path.join(temp_dir.name, file.name)
        with open(temp_file_path, "wb") as temp_file:
            temp_file.write(file.getvalue())

        #  加载器这个文件 目前在缓存里
        loader = TextLoader(temp_file_path, encoding="utf-8")
        docs.extend(loader.load())

    # 将文档分快对象进行 txt截取
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splitter = text_splitter.split_documents(docs)

    #  创建向量数据库
    # 将文本转为向量
    # embeddings = OpenAIEmbeddings(api_key=DASHSCOPE_API_KEY)
    # embeddings = Ha(api_key=DASHSCOPE_API_KEY)
    embeddings = HuggingFaceEmbeddings(
        model_name="all-MiniLM-L6-v2",
        cache_folder="./models"
    )

    vectordb = Chroma.from_documents(splitter,embeddings)

    # 创建文档检索器
    return vectordb.as_retriever()



# 配置检索器
retriever = config_retriever(uploader_file)

• 六

# 指令模版 【需要他做什么】
instructions = """您是一个设计用于查询文档来回答问题的代理。
您可以使用文档检索工具，并基于检索内容来回答问题您可能不查询文档就知道答案，但是您仍然应该查询文档来获得答案。
如果您从文档中找不到任何信息用于回答问题，则只需返回“抱歉，这个问题我还不知道。”作为答案。"""

# 基础提示词模版【类似于说明书】尽量使用英文
base_prompt_template = """
{instructions}
TOOLS:
ーーーーーー
You have access to the following tools:
{tools}
To use a tool, please use the following format:

Thought: Do I need to use a tool? Yes
Action: the action to take, should be one of [{tool_names}]
Action Input: {input}
Observation: the result of the action


When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:


Thought: Do I need to use a tool? No Final Answer: [your response here]


Begin!

Previous conversation history:
{chat_history}

New input: {input} 
{agent_scratchpad}"""

• 七 主要代码

#  创建基础提示词模版
base_prompt = PromptTemplate.from_template(base_prompt_template)

# 创建部分填充提示词模版
prompt = base_prompt.partial(
    instructions=instructions
    #instructions=instructions,
    #tools=tool.to_tool_names(),
    #tool_names=",".join(tool.to_tool_names())
)

# 创建一个 agent
agent = create_react_agent(
    llm=llm,
    tools=tools,
    prompt=prompt,
)

# 创建一个 agent 执行器
agent_executor = AgentExecutor(agent=agent,tools=tools, memory=memory, verbose=True, handle_parsing_errors="未发现数据")

user_query = st.chat_input(placeholder="请输入问题")
# 创建聊天输入框
if user_query:

    st.session_state.messages.append({"role":"user", "content":user_query})

    st.chat_message("user").write(user_query)


    # 创建一个回调
    with st.chat_message("assistant"):
        st_cb = StreamlitCallbackHandler(st.container(), expand_new_thoughts=True)
        config = {"callbacks": [st_cb]}
        response = agent_executor.invoke({"input":user_query}, config=config)
        st.session_state.messages.append({"role":"assistant", "content":response["output"]})
        st.write(response["output"])

streamlit 启动命令

# doc_chat.py  文件名称
streamlit run doc_chat.py

发布时间 2025-05-15