RAG系统开发实战

🎯 概述

RAG（Retrieval-Augmented Generation，检索增强生成）是一种结合了信息检索和文本生成的技术，能够从知识库中检索相关信息来增强大语言模型的生成能力。

📚 核心概念

RAG工作流程

1. 文档处理: 将原始文档分割成小块
2. 向量化: 使用embedding模型将文档块转换为向量
3. 存储: 将向量存储到向量数据库中
4. 检索: 根据用户问题检索相关文档块
5. 生成: 将检索到的文档和问题一起输入LLM生成答案

🛠️ 技术架构

系统组件

┌─────────────┐    ┌─────────────┐    ┌─────────────┐
│   文档处理   │───▶│   向量化    │───▶│  向量数据库  │
└─────────────┘    └─────────────┘    └─────────────┘
                                              │
┌─────────────┐    ┌─────────────┐           │
│   用户问题   │───▶│   检索系统   │───────────┘
└─────────────┘    └─────────────┘
                           │
                           ▼
                   ┌─────────────┐
                   │   LLM生成   │
                   └─────────────┘

🚀 实现步骤

1. 文档处理

import os
import json
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import DirectoryLoader

class DocumentProcessor:
    def __init__(self, chunk_size=1000, chunk_overlap=200):
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
        )
    
    def load_documents(self, directory_path):
        """加载目录中的所有文档"""
        loader = DirectoryLoader(directory_path, glob="**/*.txt")
        documents = loader.load()
        return documents
    
    def split_documents(self, documents):
        """将文档分割成小块"""
        chunks = self.text_splitter.split_documents(documents)
        return chunks
    
    def process_documents(self, directory_path):
        """完整的文档处理流程"""
        documents = self.load_documents(directory_path)
        chunks = self.split_documents(documents)
        
        # 添加元数据
        for i, chunk in enumerate(chunks):
            chunk.metadata['chunk_id'] = i
            chunk.metadata['source'] = chunk.metadata.get('source', 'unknown')
        
        return chunks

2. 向量化存储

import chromadb
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma

class VectorStore:
    def __init__(self, persist_directory="./chroma_db"):
        self.embeddings = OpenAIEmbeddings()
        self.persist_directory = persist_directory
        self.vectorstore = None
    
    def create_vectorstore(self, chunks):
        """创建向量数据库"""
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embeddings,
            persist_directory=self.persist_directory
        )
        return self.vectorstore
    
    def load_vectorstore(self):
        """加载已存在的向量数据库"""
        self.vectorstore = Chroma(
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
        return self.vectorstore
    
    def similarity_search(self, query, k=5):
        """相似性搜索"""
        if self.vectorstore is None:
            self.load_vectorstore()
        
        docs = self.vectorstore.similarity_search(query, k=k)
        return docs

3. 检索系统

from langchain.retrievers import VectorStoreRetriever
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.llms import OpenAI

class RetrievalSystem:
    def __init__(self, vectorstore, llm):
        self.vectorstore = vectorstore
        self.llm = llm
        self.retriever = self._create_retriever()
    
    def _create_retriever(self):
        """创建检索器"""
        base_retriever = VectorStoreRetriever(
            vectorstore=self.vectorstore,
            search_kwargs={"k": 5}
        )
        
        # 使用多查询检索器提高检索质量
        retriever = MultiQueryRetriever.from_llm(
            retriever=base_retriever,
            llm=self.llm
        )
        
        return retriever
    
    def retrieve_documents(self, query):
        """检索相关文档"""
        docs = self.retriever.get_relevant_documents(query)
        return docs

4. 生成系统

from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

class RAGSystem:
    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm
        self.qa_chain = self._create_qa_chain()
    
    def _create_qa_chain(self):
        """创建问答链"""
        prompt_template = """
        基于以下上下文信息回答问题。如果上下文中没有相关信息，请说"我无法从提供的信息中找到答案"。

        上下文:
        {context}

        问题: {question}

        答案:
        """
        
        PROMPT = PromptTemplate(
            template=prompt_template,
            input_variables=["context", "question"]
        )
        
        qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.retriever,
            chain_type_kwargs={"prompt": PROMPT}
        )
        
        return qa_chain
    
    def ask_question(self, question):
        """回答问题"""
        result = self.qa_chain({"query": question})
        return result["result"]

🎯 完整实现

class RAGApplication:
    def __init__(self, openai_api_key, persist_directory="./chroma_db"):
        self.openai_api_key = openai_api_key
        self.persist_directory = persist_directory
        self.document_processor = DocumentProcessor()
        self.vector_store = VectorStore(persist_directory)
        self.llm = OpenAI(openai_api_key=openai_api_key)
        self.rag_system = None
    
    def setup_knowledge_base(self, documents_directory):
        """设置知识库"""
        print("正在处理文档...")
        chunks = self.document_processor.process_documents(documents_directory)
        
        print("正在创建向量数据库...")
        self.vector_store.create_vectorstore(chunks)
        
        print("正在初始化RAG系统...")
        retrieval_system = RetrievalSystem(self.vector_store.vectorstore, self.llm)
        self.rag_system = RAGSystem(retrieval_system.retriever, self.llm)
        
        print("知识库设置完成！")
    
    def load_existing_knowledge_base(self):
        """加载已存在的知识库"""
        print("正在加载向量数据库...")
        self.vector_store.load_vectorstore()
        
        print("正在初始化RAG系统...")
        retrieval_system = RetrievalSystem(self.vector_store.vectorstore, self.llm)
        self.rag_system = RAGSystem(retrieval_system.retriever, self.llm)
        
        print("知识库加载完成！")
    
    def chat(self, question):
        """与RAG系统对话"""
        if self.rag_system is None:
            print("请先设置知识库！")
            return
        
        answer = self.rag_system.ask_question(question)
        return answer

# 使用示例
if __name__ == "__main__":
    # 初始化RAG应用
    rag_app = RAGApplication(openai_api_key="your-api-key")
    
    # 设置知识库（首次运行）
    # rag_app.setup_knowledge_base("./documents")
    
    # 加载已存在的知识库
    rag_app.load_existing_knowledge_base()
    
    # 开始对话
    while True:
        question = input("请输入问题（输入'quit'退出）: ")
        if question.lower() == 'quit':
            break
        
        answer = rag_app.chat(question)
        print(f"回答: {answer}\n")

🔧 优化技巧

1. 检索优化

def hybrid_search(self, query, alpha=0.7):
    """混合搜索：结合关键词搜索和语义搜索"""
    # 语义搜索
    semantic_docs = self.vectorstore.similarity_search(query, k=5)
    
    # 关键词搜索
    keyword_docs = self.vectorstore.similarity_search_with_score(query, k=5)
    
    # 合并结果
    combined_docs = self._combine_search_results(
        semantic_docs, keyword_docs, alpha
    )
    
    return combined_docs

def rerank_documents(self, query, documents, top_k=3):
    """重新排序文档"""
    # 使用更复杂的重排序算法
    scores = []
    for doc in documents:
        score = self._calculate_relevance_score(query, doc)
        scores.append((doc, score))
    
    # 按分数排序
    scores.sort(key=lambda x: x[1], reverse=True)
    return [doc for doc, score in scores[:top_k]]

2. 提示词优化

def create_optimized_prompt(self):
    """创建优化的提示词模板"""
    prompt_template = """
    你是一个专业的AI助手，请基于以下上下文信息回答问题。

    上下文信息:
    {context}

    用户问题: {question}

    回答要求:
    1. 基于上下文信息提供准确、详细的答案
    2. 如果上下文中没有相关信息，请明确说明
    3. 回答要结构清晰，逻辑性强
    4. 如果可能，提供具体的例子或数据支持

    答案:
    """
    
    return PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question"]
    )

3. 缓存机制

import redis
import hashlib
import json

class RAGCache:
    def __init__(self, redis_url="redis://localhost:6379"):
        self.redis_client = redis.from_url(redis_url)
        self.cache_ttl = 3600  # 1小时
    
    def get_cache_key(self, query):
        """生成缓存键"""
        return hashlib.md5(query.encode()).hexdigest()
    
    def get_cached_answer(self, query):
        """获取缓存的答案"""
        cache_key = self.get_cache_key(query)
        cached = self.redis_client.get(cache_key)
        if cached:
            return json.loads(cached)
        return None
    
    def cache_answer(self, query, answer):
        """缓存答案"""
        cache_key = self.get_cache_key(query)
        self.redis_client.setex(
            cache_key, 
            self.cache_ttl, 
            json.dumps(answer)
        )

📊 性能监控

import time
import logging
from typing import Dict, Any

class RAGMetrics:
    def __init__(self):
        self.metrics = {
            'total_queries': 0,
            'avg_response_time': 0,
            'cache_hit_rate': 0,
            'retrieval_accuracy': 0
        }
        self.logger = logging.getLogger(__name__)
    
    def log_query(self, query: str, response_time: float, 
                  cache_hit: bool, retrieved_docs: int):
        """记录查询指标"""
        self.metrics['total_queries'] += 1
        
        # 更新平均响应时间
        current_avg = self.metrics['avg_response_time']
        total_queries = self.metrics['total_queries']
        self.metrics['avg_response_time'] = (
            (current_avg * (total_queries - 1) + response_time) / total_queries
        )
        
        # 记录日志
        self.logger.info(f"Query: {query[:50]}... | "
                        f"Response time: {response_time:.2f}s | "
                        f"Cache hit: {cache_hit} | "
                        f"Retrieved docs: {retrieved_docs}")
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取性能指标"""
        return self.metrics.copy()

🎯 部署建议

1. Docker部署

FROM python:3.9-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

EXPOSE 8000

CMD ["python", "app.py"]

2. API服务

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI(title="RAG API", version="1.0.0")

class QueryRequest(BaseModel):
    question: str
    max_tokens: int = 500

class QueryResponse(BaseModel):
    answer: str
    sources: list
    response_time: float

@app.post("/query", response_model=QueryResponse)
async def query_rag(request: QueryRequest):
    try:
        start_time = time.time()
        answer = rag_app.chat(request.question)
        response_time = time.time() - start_time
        
        return QueryResponse(
            answer=answer,
            sources=[],  # 可以添加来源文档
            response_time=response_time
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

🎉 总结

RAG系统开发涉及多个技术环节，从文档处理到向量检索，再到生成优化。通过合理的架构设计和优化策略，可以构建出高质量的RAG应用。

关键成功因素：

1. 高质量的文档预处理
2. 合适的embedding模型选择
3. 有效的检索策略
4. 优化的提示词设计
5. 完善的监控和缓存机制

---

RAG让AI更智能！ 🤖