第 12 章 企业 RAG 知识库与 DeepSeek 源码级集成

12.1 向量库(Milvus/FAISS)推理层深度融合改造

12.1.1 RAG 架构概述

RAG(Retrieval-Augmented Generation)是一种将检索与生成相结合的技术,通过从知识库中检索相关信息来增强模型的生成能力。

RAG 流程:

  1. 文档索引:将文档转换为向量并存储到向量库
  2. 查询检索:将用户查询转换为向量,在向量库中检索相似文档
  3. 上下文构建:将检索到的文档构建为上下文
  4. 生成回答:基于上下文和查询生成回答

12.1.2 Milvus 集成

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
import numpy as np

class MilvusRetriever:
    def __init__(self, collection_name="deepseek_rag", dim=768):
        self.collection_name = collection_name
        self.dim = dim
        
        connections.connect("default", host="localhost", port="19530")
        
        if not utility.has_collection(self.collection_name):
            self._create_collection()
        
        self.collection = Collection(self.collection_name)
        self.collection.load()
    
    def _create_collection(self):
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.dim),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="metadata", dtype=DataType.JSON)
        ]
        
        schema = CollectionSchema(fields=fields, description="DeepSeek RAG Knowledge Base")
        
        Collection(name=self.collection_name, schema=schema)
    
    def add_documents(self, documents, embeddings):
        entities = [
            embeddings,
            [doc["text"] for doc in documents],
            [doc.get("metadata", {}) for doc in documents]
        ]
        
        self.collection.insert(entities)
        self.collection.flush()
    
    def retrieve(self, query_embedding, top_k=5):
        search_params = {
            "metric_type": "COSINE",
            "params": {"ef": 128}
        }
        
        results = self.collection.search(
            data=[query_embedding],
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            output_fields=["text", "metadata"]
        )
        
        retrieved_docs = []
        for hit in results[0]:
            retrieved_docs.append({
                "text": hit.entity.get("text"),
                "metadata": hit.entity.get("metadata"),
                "score": hit.score
            })
        
        return retrieved_docs

12.1.3 FAISS 集成

import faiss
import numpy as np

class FAISSRetriever:
    def __init__(self, index_path=None, dim=768):
        self.dim = dim
        self.index = faiss.IndexFlatIP(dim)
        
        if index_path and os.path.exists(index_path):
            self.index = faiss.read_index(index_path)
        
        self.documents = []
    
    def add_documents(self, documents, embeddings):
        self.documents.extend(documents)
        
        embeddings_np = np.array(embeddings, dtype=np.float32)
        self.index.add(embeddings_np)
    
    def retrieve(self, query_embedding, top_k=5):
        query_np = np.array([query_embedding], dtype=np.float32)
        
        distances, indices = self.index.search(query_np, top_k)
        
        retrieved_docs = []
        for i, idx in enumerate(indices[0]):
            if idx < len(self.documents):
                retrieved_docs.append({
                    "text": self.documents[idx]["text"],
                    "metadata": self.documents[idx].get("metadata", {}),
                    "score": float(distances[0][i])
                })
        
        return retrieved_docs
    
    def save_index(self, index_path):
        faiss.write_index(self.index, index_path)

12.1.4 向量库对比

特性MilvusFAISS
分布式支持支持不支持
持久化存储支持需手动管理
查询性能极高
部署复杂度
适用场景大规模生产小规模实验

12.2 检索增强生成缓存优化源码

12.2.1 RAG 缓存设计

from functools import lru_cache
from datetime import datetime, timedelta

class RAGCache:
    def __init__(self, max_size=10000, ttl_hours=24):
        self.cache = {}
        self.max_size = max_size
        self.ttl = timedelta(hours=ttl_hours)
    
    def get(self, query):
        if query in self.cache:
            entry = self.cache[query]
            
            if datetime.now() - entry["timestamp"] < self.ttl:
                return entry["result"]
            
            del self.cache[query]
        
        return None
    
    def set(self, query, result):
        if len(self.cache) >= self.max_size:
            oldest_key = min(self.cache.keys(), key=lambda k: self.cache[k]["timestamp"])
            del self.cache[oldest_key]
        
        self.cache[query] = {
            "result": result,
            "timestamp": datetime.now()
        }
    
    def clear(self):
        self.cache.clear()
    
    def size(self):
        return len(self.cache)

12.2.2 RAG 流程优化

class OptimizedRAG:
    def __init__(self, retriever, model, tokenizer, embedding_model):
        self.retriever = retriever
        self.model = model
        self.tokenizer = tokenizer
        self.embedding_model = embedding_model
        self.cache = RAGCache()
    
    def generate(self, query, max_tokens=512, top_k=5):
        cached_result = self.cache.get(query)
        
        if cached_result:
            return cached_result
        
        query_embedding = self.embedding_model.encode(query)
        
        retrieved_docs = self.retriever.retrieve(query_embedding, top_k=top_k)
        
        context = "
".join([doc["text"] for doc in retrieved_docs])
        
        prompt = f"""基于以下上下文回答问题:

上下文:
{context}

问题:{query}

请根据上下文内容回答问题,如果上下文没有相关信息,请说明。"""
        
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").cuda()
        
        with torch.no_grad():
            output = self.model.generate(input_ids, max_new_tokens=max_tokens)
        
        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
        
        self.cache.set(query, response)
        
        return response

12.2.3 缓存命中率统计

class CacheStats:
    def __init__(self):
        self.hits = 0
        self.misses = 0
        self.total_queries = 0
    
    def record_hit(self):
        self.hits += 1
        self.total_queries += 1
    
    def record_miss(self):
        self.misses += 1
        self.total_queries += 1
    
    def get_hit_rate(self):
        if self.total_queries == 0:
            return 0.0
        
        return self.hits / self.total_queries
    
    def reset(self):
        self.hits = 0
        self.misses = 0
        self.total_queries = 0

12.3 私有文档解析、分段重排定制代码

12.3.1 文档解析

import PyPDF2
import docx
from pptx import Presentation
import markdown

class DocumentParser:
    def __init__(self):
        self.parsers = {
            ".pdf": self._parse_pdf,
            ".docx": self._parse_docx,
            ".pptx": self._parse_pptx,
            ".md": self._parse_md,
            ".txt": self._parse_txt
        }
    
    def parse(self, file_path):
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext not in self.parsers:
            raise ValueError(f"不支持的文件格式: {ext}")
        
        return self.parsers[ext](file_path)
    
    def _parse_pdf(self, file_path):
        text = ""
        
        with open(file_path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            
            for page in reader.pages:
                text += page.extract_text() + "
"
        
        return text
    
    def _parse_docx(self, file_path):
        doc = docx.Document(file_path)
        text = ""
        
        for paragraph in doc.paragraphs:
            text += paragraph.text + "
"
        
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    text += cell.text + " "
                text += "
"
        
        return text
    
    def _parse_pptx(self, file_path):
        prs = Presentation(file_path)
        text = ""
        
        for slide in prs.slides:
            for shape in slide.shapes:
                if hasattr(shape, "text"):
                    text += shape.text + "
"
        
        return text
    
    def _parse_md(self, file_path):
        with open(file_path, "r", encoding="utf-8") as f:
            md_text = f.read()
        
        html = markdown.markdown(md_text)
        text = "".join(BeautifulSoup(html, "html.parser").findAll(text=True))
        
        return text
    
    def _parse_txt(self, file_path):
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()

12.3.2 文档分段

class DocumentChunker:
    def __init__(self, chunk_size=512, chunk_overlap=50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
    
    def chunk(self, text):
        chunks = []
        
        words = text.split()
        total_words = len(words)
        
        start = 0
        while start < total_words:
            end = start + self.chunk_size
            
            if end < total_words:
                chunk_words = words[start:end]
            else:
                chunk_words = words[start:]
            
            chunk_text = " ".join(chunk_words)
            chunks.append(chunk_text)
            
            start += self.chunk_size - self.chunk_overlap
        
        return chunks
    
    def chunk_with_overlap(self, text):
        chunks = []
        
        sentences = text.split("。")
        current_chunk = ""
        current_length = 0
        
        for sentence in sentences:
            sentence_length = len(sentence)
            
            if current_length + sentence_length <= self.chunk_size:
                current_chunk += sentence + "。"
                current_length += sentence_length
            else:
                if current_chunk:
                    chunks.append(current_chunk)
                
                current_chunk = sentence + "。"
                current_length = sentence_length
        
        if current_chunk:
            chunks.append(current_chunk)
        
        return chunks

12.3.3 文档重排

class DocumentReranker:
    def __init__(self, rerank_model):
        self.rerank_model = rerank_model
    
    def rerank(self, query, documents):
        texts = [doc["text"] for doc in documents]
        
        scores = self.rerank_model.compute_score(query, texts)
        
        ranked_docs = sorted(
            zip(documents, scores),
            key=lambda x: x[1],
            reverse=True
        )
        
        return [doc for doc, _ in ranked_docs]

12.4 幻觉抑制底层逻辑修改

12.4.1 幻觉检测

class HallucinationDetector:
    def __init__(self, embedding_model, threshold=0.7):
        self.embedding_model = embedding_model
        self.threshold = threshold
    
    def detect(self, generated_text, context_documents):
        context_embeddings = [
            self.embedding_model.encode(doc["text"])
            for doc in context_documents
        ]
        
        generated_embedding = self.embedding_model.encode(generated_text)
        
        similarities = []
        for ctx_emb in context_embeddings:
            similarity = self._cosine_similarity(generated_embedding, ctx_emb)
            similarities.append(similarity)
        
        avg_similarity = sum(similarities) / len(similarities)
        
        is_hallucination = avg_similarity < self.threshold
        
        return {
            "is_hallucination": is_hallucination,
            "similarity_score": avg_similarity,
            "threshold": self.threshold
        }
    
    def _cosine_similarity(self, vec1, vec2):
        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        
        return dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0

12.4.2 幻觉抑制

class HallucinationSuppressor:
    def __init__(self, detector, max_retries=3):
        self.detector = detector
        self.max_retries = max_retries
    
    def suppress(self, query, context_documents, generate_fn):
        for attempt in range(self.max_retries):
            response = generate_fn(query, context_documents)
            
            detection = self.detector.detect(response, context_documents)
            
            if not detection["is_hallucination"]:
                return response, detection
            
            print(f"检测到幻觉,尝试重新生成 (第 {attempt + 1} 次)")
        
        return response, detection

12.4.3 防幻觉提示词

class AntiHallucinationPrompt:
    def __init__(self):
        self.system_prompt = """你是一个专业的问答助手,必须严格基于提供的上下文回答问题。

规则:
1. 只使用上下文提供的信息回答问题
2. 如果上下文没有相关信息,明确说明"根据提供的上下文,无法回答该问题"
3. 不要编造任何信息
4. 如果不确定答案,说明"无法确定"

请严格遵守以上规则。"""
    
    def build_prompt(self, query, context_documents):
        context = "

".join([doc["text"] for doc in context_documents])
        
        prompt = f"""{self.system_prompt}

上下文:
{context}

问题:{query}

回答:"""
        
        return prompt

本章小结:

本章详细介绍了企业 RAG 知识库与 DeepSeek 的源码级集成,包括向量库集成、缓存优化、文档解析分段和幻觉抑制。这些技术为企业构建智能问答系统提供了完整的解决方案。
更多资讯:lxb20110121

Logo

欢迎加入DeepSeek 技术社区。在这里,你可以找到志同道合的朋友,共同探索AI技术的奥秘。

更多推荐