work_02/rag_practice.py

@@ -0,0 +1,205 @@
+"""
+一. 构建知识库
+    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)