""" 一. 构建知识库 1. 收集资料(无优化空间,但为质量的前提) 2. 解析并清洗数据 ①清洗噪音 ②清洗多余无意义字符 3. 分块(决定7成左右RAG质量) ①按固定字符切分 ②按递归符号切分(递归符号:需人工分析,少量可官网AI总结) ③按文档块标识符切分(一般固定文档格式类型,如markdown) ④按语义模块切分(费钱,使用embadding模型) ⑤按LLM切分(烧钱,不考虑钱直接冲) 4. 向量化(智普embadding模型, 余弦相似度) 5. 向量数据入库(测试chroma,主流milvus) 二. 用户基于知识库提问 1. 用户提问 2. 用户问题向量化,到库中检索,top_k返回前n条完整结果 3. 将结果封装进提示词 4. 请求LLM返回最终结果 """ from dotenv import load_dotenv from langchain_openai import ChatOpenAI import os, re from fastapi import FastAPI, UploadFile, File, HTTPException from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader import tempfile from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_community.vectorstores import Chroma from langchain_community.embeddings import ZhipuAIEmbeddings from langchain_core.prompts import PromptTemplate from zai import ZhipuAiClient load_dotenv(override=True) deepseek_base_url = os.getenv("DEEPSEEK_BASE_URL") deepseek_base_key = os.getenv("DEEPSEEK_BASE_KEY") deepseek_base_name = os.getenv("DEEPSEEK_BASE_NAME") # ==============================一. 构建知识库================================================== # ==============================1. 收集资料================================================== class KnowledgeBaseBuilder: def __init__(self, pdf_dir="./resources", persist_dir="./chroma_db"): self.pdf_dir = pdf_dir self.persist_dir = persist_dir self.embeddings = ZhipuAIEmbeddings( model="embedding-3", # 或者 "embedding-3" api_key=os.getenv("ZHIPU_API_KEY") ) self.documents = [] def load_pdfs(self): """加载PDF文件""" loader = DirectoryLoader( self.pdf_dir, glob="**/*.pdf", loader_cls=PyPDFLoader, show_progress=True ) self.documents = loader.load() print(f"加载了 {len(self.documents)} 个文档") return self.documents # ==============================2. 解析并清洗数据================================================== def clean_documents(self, docs): """清洗文档数据""" cleaned_docs = [] empty_count = 0 for i, doc in enumerate(docs): text = doc.page_content # 检查文本是否为空 if not text or len(text.strip()) == 0: empty_count += 1 print(f"⚠️ 文档 {i+1} 内容为空,跳过") continue # ① 清洗噪音:移除页眉页脚、水印等 text = re.sub(r'第\s*\d+\s*页\s*/\s*共\s*\d+\s*页', '', text) text = re.sub(r'Copyright.*?\n', '', text, flags=re.IGNORECASE) text = re.sub(r'机密|内部资料|仅供内部使用', '', text) # ② 去除无意义字符 text = re.sub(r'\n\s*\n+', '\n\n', text) # 合并多个空行 text = re.sub(r'[ \t]+', ' ', text) # 合并多个空格 # 保留中英文、数字和常用标点 text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\.\,\,\。\!\?\:\;\(\)\n]', ' ', text) text = re.sub(r'\s+', ' ', text).strip() # 清洗后再次检查 if not text or len(text) < 10: empty_count += 1 print(f"⚠️ 文档 {i+1} 清洗后内容过短,跳过") continue doc.page_content = text cleaned_docs.append(doc) print(f"✅ 清洗完成,有效文档: {len(cleaned_docs)}, 跳过: {empty_count}") if not cleaned_docs: raise ValueError("清洗后没有有效的文档内容") return cleaned_docs # ==============================3. 分块================================================== def chunk_documents(self, docs, strategy="recursive"): """ 多种分块策略 strategy: fixed, recursive, semantic, markdown 当前实现递归分块策略 """ if strategy == "recursive": splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=50, separators=[ "\n\n", # 段落 "\n", # 行 "。", # 中文句号 "!", # 感叹号 "?", # 问号 ";", # 分号 ",", # 逗号 ".", # 英文句号 "!", "?", " " ], length_function=len, keep_separator=True ) chunked_docs = splitter.split_documents(docs) return chunked_docs # ==============================4. 向量化================================================== # ==============================5. 向量数据入库================================================== def build_vectorstore(self, chunks, collection_name="knowledge_base"): """向量化并存入Chroma数据库""" # 使用智谱Embedding模型 vectorstore = Chroma.from_documents( documents=chunks, embedding=self.embeddings, persist_directory=self.persist_dir, collection_name=collection_name ) vectorstore.persist() print(f"向量数据库构建完成,存储于: {self.persist_dir}") return vectorstore # ==============================二. 用户基于知识库提问================================================== # ==============================1. 用户提问================================================== # ==============================2. 用户问题向量化,到库中检索,top_k返回前n条完整结果================================================== # ==============================3. 将结果封装进提示词================================================== # ==============================4. 请求LLM返回最终结果================================================== llm = ChatOpenAI( base_url = deepseek_base_url, api_key = deepseek_base_key, model= deepseek_base_name ) if __name__ == "__main__": knowlege = KnowledgeBaseBuilder() docs = knowlege.load_pdfs() clean_docs = knowlege.clean_documents(docs) chunks = knowlege.chunk_documents(clean_docs) vectorstore = knowlege.build_vectorstore(chunks) # vectorstore = Chroma( # persist_directory="./chroma_db", # 持久化目录 # embedding_function=knowlege.embeddings, # collection_name="knowledge_base" # 集合名称(与构建时一致) # ) user_question = '遥控器如何操作?' retrieved_docs = vectorstore.similarity_search(user_question, k=3) context = "\n\n".join([doc.page_content for doc in retrieved_docs]) prompt_template = """ 你是一个专业的AI助手,请基于以下参考资料回答用户的问题。 【参考资料】: {context} 【用户问题】: {question} 要求: 1. 严格基于参考资料回答,不要编造信息 2. 如果参考资料中没有相关信息,请明确告知 3. 回答要简洁、准确、有条理 回答: """ prompt = PromptTemplate( template=prompt_template, input_variables=["context", "question"] ) # 4. 请求LLM formatted_prompt = prompt.format(context=context, question=user_question) response = llm.invoke(formatted_prompt) print(response)