from langchain.document_loaders import PyPDFLoader, DirectoryLoaderfrom langchain.text_splitter import RecursiveCharacterTextSplitterfrom langchain.embeddings import HuggingFaceEmbeddings  from langchain.vectorstores import Chroma from langchain_community.llms import Ollamafrom langchain.chains import RetrievalQA # 1. 配置路径 DOC_FOLDER = ”./docs”# 存放你的PDF/TXT文档VECTOR_DB_PATH = ”./chroma_db”EMBED_MODEL = ”all-MiniLM-L6-v2”LLM_MODEL = ”qwen:7b”# 2. 加载所有文档loader = DirectoryLoader(DOC_FOLDER, glob=”**/*.pdf”, loader_cls=PyPDFLoader) documents = loader.load()# 3. 文本切块 text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=80, separators=[”\n\n”, ”\n”, ”。”, ”，”, ” ”] )texts = text_splitter.split_documents(documents)# 4. 本地离线Embedding embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)# 5. 构建/加载向量库vectordb = Chroma.from_documents( documents=texts, embedding=embeddings, persist_directory=VECTOR_DB_PATH)vectordb.persist()# 6. 加载本地大模型Ollama（提前执行 ollama pull qwen:7b）  llm = Ollama(model=LLM_MODEL)# 7. 检索链：召回3段最相关内容retriever = vectordb.as_retriever(search_kwargs={”k”: 3})qa_chain = RetrievalQA.from_chain_type( llm=llm, chain_type=”stuff”, retriever=retriever, return_source_documents=True  # 返回参考文档片段   )# 对话循环if __name__ == ”__main__”:   print(”本地RAG知识库就绪，输入exit退出”) while True: query = input(”提问：”)  if query == ”exit”: break res = qa_chain.invoke({”query”: query}) print(”\n回答：”, res[”result”]) print(”\n参考来源片段：”) for idx, doc inenumerate(res[”source_documents”]):  print(f”{idx+1}. {doc.page_content[:200]}...”)

你是基于本地知识库的问答助手，仅使用下方提供的参考资料回答用户问题。   如果资料中没有对应信息，直接回复“知识库未收录相关内容”，禁止编造、猜测外部知识。     参考资料：     {context}     用户问题：{question}

pip install streamlit langchain langchain-community chromadb sentence-transformers pypdf2 ollama python-multipart

# 大模型ollama pull qwen:7b# 嵌入模型代码自动下载，无需手动操作

rag_streamlit/  ├── streamlit_rag.py #主程序  ├── docs_upload/ # 自动生成，存放上传文件  └── chroma_db/ # 向量库持久化目录

import os    import streamlit as st from PyPDF2 import PdfReaderfrom langchain.text_splitter import RecursiveCharacterTextSplitterfrom langchain.embeddings import HuggingFaceEmbeddingsfrom langchain.vectorstores import Chromafrom langchain_community.llms import Ollamafrom langchain.chains import RetrievalQAfrom langchain.prompts import PromptTemplate# ===================== 全局配置 =====================UPLOAD_FOLDER = ”./docs_upload”CHROMA_DB_DIR = ”./chroma_db”EMBED_MODEL_NAME = ”all-MiniLM-L6-v2”LLM_MODEL_NAME = ”qwen:7b”# 创建文件夹os.makedirs(UPLOAD_FOLDER, exist_ok=True)os.makedirs(CHROMA_DB_DIR, exist_ok=True)# 页面基础设置st.set_page_config(page_title=”个人本地RAG知识库”, layout=”wide”)st.title(”📚 私有化本地RAG知识库”)# ===================== 初始化模型与向量库 =====================@st.cache_resource(show_spinner=”加载嵌入模型...”)def get_embedding_model(): return HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)@st.cache_resource(show_spinner=”加载本地大模型...”)def get_llm(): return Ollama(model=LLM_MODEL_NAME, temperature=0.1)def load_vector_db(): embeddings = get_embedding_model()  db = Chroma(persist_directory=CHROMA_DB_DIR, embedding_function=embeddings)   return db # 自定义Prompt，禁止模型编造内容CUSTOM_PROMPT_TPL = ”””你是专属本地知识库问答助手，严格遵守以下规则：  1. 仅使用下方【参考资料】中的内容回答用户问题；   2. 如果参考资料没有相关信息，直接回复：「知识库未收录相关内容，无法解答」；  3. 禁止凭空猜测、补充外部互联网知识；   4. 回答简洁清晰，分点说明。  【参考资料】{context}用户问题：{question}  ”””PROMPT = PromptTemplate( template=CUSTOM_PROMPT_TPL,  input_variables=[”context”, ”question”])# ===================== 文件处理工具 =====================def extract_pdf_text(file_path): text = ”” pdf_reader = PdfReader(file_path) for page in pdf_reader.pages: page_text = page.extract_text() if page_text:  text += page_text + ”\n” return textdef extract_txt_text(file_path): with open(file_path, ”r”, encoding=”utf-8”) as f:  return f.read()def split_text(text): splitter = RecursiveCharacterTextSplitter( chunk_size=600, chunk_overlap=80, separators=[”\n\n”, ”\n”, ”。”, ”，”, ” ”]  ) return splitter.split_text(text)def add_text_to_vector_db(text_list, source_name): db = load_vector_db()   metadatas = [{”source”: source_name} for _ in text_list] db.add_texts(texts=text_list, metadatas=metadatas) db.persist() # ===================== 侧边栏：文件上传 & 库管理 =====================with st.sidebar: st.header(”📂 文档上传管理”)  uploaded_files = st.file_uploader(”上传 PDF / TXT 文件”, type=[”pdf”, ”txt”], accept_multiple_files=True)  if uploaded_files: for file in uploaded_files: save_path = os.path.join(UPLOAD_FOLDER, file.name)# 保存本地文件 with open(save_path, ”wb”) as f: f.write(file.read())  st.success(f”已保存：{file.name}”)# 提取文本 if file.name.endswith(”.pdf”): content = extract_pdf_text(save_path) else: content = extract_txt_text(save_path)# 切片入库 chunks = split_text(content) add_text_to_vector_db(chunks, file.name) st.info(f”{file.name} 已完成向量化入库，切片数量：{len(chunks)}”) st.divider() st.header(”🗑️ 向量库操作”) if st.button(”清空全部知识库”, type=”secondary”): db = load_vector_db() db.delete_collection() st.cache_resource.clear() st.rerun()# ===================== 主页面对话区域 =====================st.subheader(”💬 基于私有文档问答”)user_query = st.chat_input(”输入你的问题...”)if user_query: with st.spinner(”检索知识库 + 生成回答中...”):  db = load_vector_db() llm = get_llm()  # 检索器，返回top3相关片段 retriever = db.as_retriever(search_kwargs={”k”: 3}) qa_chain = RetrievalQA.from_chain_type( llm=llm, chain_type=”stuff”, retriever=retriever, return_source_documents=True, chain_type_kwargs={”prompt”: PROMPT} ) result = qa_chain.invoke({”query”: user_query})# 输出回答 st.markdown(”### 回答”) st.write(result[”result”])# 输出引用来源 st.markdown(”### 📄 参考文档片段”) source_docs = result[”source_documents”] if not source_docs: st.write(”无匹配文档”) else: seen_source = set() for idx, doc in enumerate(source_docs): src = doc.metadata[”source”] if src not in seen_source: seen_source.add(src) st.expander(f”来源文件：{src}”).write(doc.page_content)

streamlit run streamlit_rag.py