AI大模型开发 从 0 到企业级:RAG检索增强全流程实现 + 本地私有知识库搭建实战

langchain==0.1.20langchain-community==0.0.38langchain-chroma==0.1.2sentence-transformers==2.7.0pypdf==4.2.0python-docx==1.1.2bitsandbytes==0.43.1transformers==4.41.2torch==2.3.0psycopg2-binary==2.9.9pgvector==0.2.5fastapi==0.110.2uvicorn==0.29.0python-jose[cryptography]==3.3.0passlib[bcrypt]==1.7.4

pip install -r requirements.txt

import osimport refrom typing import Listfrom langchain.schema import Documentfrom pypdf import PdfReaderfrom docx import Document as DocxDocumentdef load_single_pdf(file_path: str) -> List[Document]:    pdf_reader = PdfReader(file_path)    text = ""    for page in pdf_reader.pages:        text += page.extract_text() + "\n"    return [Document(page_content=text, metadata={"source": file_path, "file_type": "pdf"})]def load_single_docx(file_path: str) -> List[Document]:    doc = DocxDocument(file_path)    text = "\n".join([para.text for para in doc.paragraphs])    return [Document(page_content=text, metadata={"source": file_path, "file_type": "docx"})]def load_text_file(file_path: str) -> List[Document]:    with open(file_path, "r", encoding="utf-8") as f:        text = f.read()    ext = os.path.splitext(file_path)[1]    return [Document(page_content=text, metadata={"source": file_path, "file_type": ext})]def load_dir_documents(dir_path: str) -> List[Document]:    documents = []    for root, _, files in os.walk(dir_path):        for file in files:            file_path = os.path.join(root, file)            ext = file.lower().split(".")[-1]            if ext == "pdf":                documents.extend(load_single_pdf(file_path))            elif ext in ["docx", "doc"]:                documents.extend(load_single_docx(file_path))            elif ext in ["txt", "md"]:                documents.extend(load_text_file(file_path))    return documentsdef preprocess_documents(documents: List[Document]) -> List[Document]:    processed = []    seen_content = set()    for doc in documents:        content = doc.page_content.strip()        content = "\n".join([line.strip() for line in content.splitlines() if line.strip()])        content = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\xff]', '', content)        if content and content not in seen_content:            seen_content.add(content)            doc.page_content = content            processed.append(doc)    return processed

from langchain.text_splitter import RecursiveCharacterTextSplitterdef split_documents(documents: List[Document], chunk_size: int = 512, chunk_overlap: int = 50) -> List[Document]:    text_splitter = RecursiveCharacterTextSplitter(        chunk_size=chunk_size,        chunk_overlap=chunk_overlap,        add_start_index=True,        separators=["\n\n", "\n", "。", "，", " ", ""]    )    return text_splitter.split_documents(documents)

import torchfrom langchain_community.embeddings import HuggingFaceBgeEmbeddingsmodel_name = "BAAI/bge-base-zh-v1.5"model_kwargs = {"device": "cuda" if torch.cuda.is_available() else "cpu", "load_in_4bit": True}encode_kwargs = {"normalize_embeddings": True}embeddings = HuggingFaceBgeEmbeddings(    model_name=model_name,    model_kwargs=model_kwargs,    encode_kwargs=encode_kwargs)

from langchain_chroma import Chroma# 初始化向量库persist_directory = "./chroma_db"vector_db = Chroma(    persist_directory=persist_directory,    embedding_function=embeddings,    collection_name="private_knowledge")# 插入分块后的文档split_docs = split_documents(preprocess_documents(load_dir_documents("./my_docs")))vector_db.add_documents(split_docs)# 持久化存储vector_db.persist()

version: '3.8'services:  postgres:    image: pgvector/pgvector:pg16    container_name: pgvector-rag    environment:      POSTGRES_USER: rag_user      POSTGRES_PASSWORD: rag_2024_secure_private      POSTGRES_DB: rag_db    ports:      - "5432:5432"    volumes:      - ./pg_data:/var/lib/postgresql/data    restart: unless-stopped

docker-compose up -d

from langchain_community.vectorstores import PGVectorconnection_string = "postgresql+psycopg2://rag_user:rag_2024_secure_private@localhost:5432/rag_db"collection_name = "private_knowledge"vector_db = PGVector(    embedding=embeddings,    collection_name=collection_name,    connection_string=connection_string,    use_jsonb=True,)split_docs = split_documents(preprocess_documents(load_dir_documents("./my_docs")))vector_db.add_documents(split_docs)

from langchain.retrievers import BM25Retriever, EnsembleRetrieverfrom langchain_community.cross_encoders import HuggingFaceCrossEncoderfrom langchain.retrievers import ContextualCompressionRetrieverfrom langchain.retrievers.document_compressors import CrossEncoderRerank# 初始化BM25稀疏召回bm25_retriever = BM25Retriever.from_documents(split_docs)bm25_retriever.k = 10# 初始化稠密向量召回vector_retriever = vector_db.as_retriever(search_kwargs={"k": 10})# 融合两路召回结果ensemble_retriever = EnsembleRetriever(    retrievers=[bm25_retriever, vector_retriever],    weights=[0.4, 0.6])# BGE重排序模型，筛选最相关的Top5rerank_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")compressor = CrossEncoderRerank(model=rerank_model, top_n=5)compression_retriever = ContextualCompressionRetriever(    base_retriever=ensemble_retriever,    base_compressor=compressor)

from langchain_community.llms import HuggingFacePipelinefrom langchain.chains import RetrievalQAfrom transformers import AutoTokenizer, AutoModelForCausalLM, pipelinefrom langchain.prompts import PromptTemplate# 私有知识库专属Prompt模板template = """你是一个专业的企业内部助手，你需要根据提供的上下文信息回答用户的问题，如果你不知道答案，就直接说你不知道，不要编造答案。上下文信息如下:{context}用户问题: {question}回答:"""prompt = PromptTemplate.from_template(template)# 加载4bit量化Llama3-8B-Instructmodel_name = "NousResearch/Meta-Llama-3-8B-Instruct"tokenizer = AutoTokenizer.from_pretrained(model_name)model = AutoModelForCausalLM.from_pretrained(    model_name,    load_in_4bit=True,    device_map="auto")pipe = pipeline(    "text-generation",    model=model,    tokenizer=tokenizer,    max_new_tokens=512,    temperature=0.1,    top_p=0.9,    repetition_penalty=1.1)llm = HuggingFacePipeline(pipeline=pipe)# 构建RAG问答链qa_chain = RetrievalQA.from_chain_type(    llm=llm,    retriever=compression_retriever,    chain_type_kwargs={"prompt": prompt},    return_source_documents=True)# 测试问答question = "请问公司的年假申请规则是什么？"result = qa_chain.invoke({"query": question})print("回答：", result["result"])print("参考文档：", [doc.metadata["source"] for doc in result["source_documents"]])

import loggingimport refrom fastapi import FastAPI, Depends, HTTPException, statusfrom jose import JWTError, jwtfrom pydantic import BaseModelfrom typing import Optional# 初始化日志审计logging.basicConfig(filename='./rag_audit.log', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')app = FastAPI(title="企业私有知识库RAG API", version="1.0")SECRET_KEY = "替换为你自己的生产级密钥"ALGORITHM = "HS256"class QueryRequest(BaseModel):    question: str    token: str# JWT认证依赖async def get_current_user(token: str) -> Optional[int]:    credentials_exception = HTTPException(        status_code=status.HTTP_401_UNAUTHORIZED,        detail="无效的访问凭证",    )    try:        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])        user_id: int = payload.get("user_id")        if user_id is None:            raise credentials_exception    except JWTError:        raise credentials_exception    return user_id# 敏感内容检查SENSITIVE_WORDS = ["替换为你的敏感词列表"]def sensitive_check(text: str) -> bool:    for word in SENSITIVE_WORDS:        if word in text:            return False    return True# PII脱敏def pii_mask(text: str) -> str:    text = re.sub(r'(1[3-9]\d)\d{4}(\d{4})', r'\1****\2', text)    text = re.sub(r'(\d{6})\d{8}(\d{4})', r'\1********\2', text)    return text# 问答接口@app.post("/api/rag/query")async def query_rag(request: QueryRequest, user_id: int = Depends(get_current_user)):    logging.info(f"用户ID: {user_id}, 请求问题: {request.question}")    # 输入敏感检查    if not sensitive_check(request.question):        logging.warning(f"用户ID: {user_id}, 请求包含敏感内容: {request.question}")        raise HTTPException(status_code=400, detail="请求包含违规内容")    # 权限过滤：只返回当前用户有权限的文档    result = qa_chain.invoke({"query": request.question, "filter": {"accessible_user": [user_id, -1]}})    # 输出脱敏    answer = pii_mask(result["result"])    # 输出敏感检查    if not sensitive_check(answer):        logging.warning(f"用户ID: {user_id}, 生成回答包含敏感内容, 问题: {request.question}")        return {"answer": "回答包含敏感内容，无法展示", "sources": []}    sources = [doc.metadata["source"] for doc in result["source_documents"]]    return {"answer": answer, "sources": sources}