feat: 新增Rag

work_02/rag_pratice.py

@@ -0,0 +1,79 @@
+import os
+import re
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter 
+from langchain_community.embeddings import DashScopeEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.chat_models import ChatTongyi
+from langchain_core.output_parsers import StrOutputParser
+from dotenv import load_dotenv
+load_dotenv()
+
+# 1.文档收集
+pdfloader = PyMuPDFLoader('./resources/bottle_info.pdf')
+pages = pdfloader.load()
+# 2.数据清洗
+def clean_pdf_text(text: str) -> str:
+  if not text:
+    return ""
+  text = re.sub(r'第\s*\d+\s*页', '', text)
+  text = re.sub(r'^\d+\s*$', '', text, flags=re.MULTILINE)
+  text = re.sub(r'[ ]{2,}', ' ', text)
+  text = re.sub(r'(?<![。！？；.!?])\n', ' ', text)
+  text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s。！？，、；：“”‘’（）《》]', '', text)
+  text = re.sub(r'\n{2,}', '\n', text)
+  return text.strip()
+
+for page in pages:
+  page.page_content = clean_pdf_text(page.page_content)
+
+pages = [p for p in pages if p.page_content.strip()]
+
+# 3.分块
+splitter = RecursiveCharacterTextSplitter(
+  separators=["\n\n", "\n", "。", "！", "？", " ", ""],
+  chunk_size = 50,
+  chunk_overlap = 10,
+  length_function = len
+)
+
+docs = splitter.split_documents(pages)
+# 4.存入向量库
+embedding_model = DashScopeEmbeddings(
+    model="text-embedding-v3",
+    dashscope_api_key= os.getenv("QW_KEY")
+)
+vectorstore = Chroma.from_documents(
+  documents = docs,
+  embedding = embedding_model,
+  persist_directory="./test_bottle_knowledge_db"
+)
+# 5.创建检索器
+retriver = vectorstore.as_retriever(search_kwargs={"k": 3})
+# 6.提问
+query = "杯型的分类有哪些？"
+relevant_docs = retriver.invoke(query)
+# 7.生成回答
+context = "\n\n---\n\n".join([d.page_content for d in relevant_docs])
+
+prompt = ChatPromptTemplate.from_template("""
+你是一个专业的知识库助手。请根据以下上下文回答问题。
+
+**规则：**
+- 只基于提供的上下文回答，不要编造
+- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」
+- 回答要简洁直接，引用原文时用引号
+
+**上下文：**
+{context}
+
+**问题：**
+{question}
+""")
+
+llm = ChatTongyi(model="qwen-plus", dashscope_api_key=os.getenv("QW_KEY"))
+chain = prompt | llm | StrOutputParser()
+
+answer = chain.invoke({"context": context, "question": query})
+print(answer)