📚 RAG完整教程

# 生产级RAG完整代码
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core.postprocessor import CohereRerank, SimilarityPostprocessor
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
import qdrant_client
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ProductionRAG:
    def __init__(self):
        # 配置Embedding
        Settings.embed_model = HuggingFaceEmbedding(
            model_name="BAAI/bge-large-zh-v1.5",
            cache_folder="./models"
        )
        
        # 配置LLM
        Settings.llm = OpenAI(
            model="gpt-4o-mini",
            temperature=0.1,
            max_tokens=1000
        )
        
        # 向量数据库
        self.client = qdrant_client.QdrantClient(
            host="localhost",
            port=6333,
            timeout=60
        )
        
        self.vector_store = QdrantVectorStore(
            client=self.client,
            collection_name="production_kb"
        )
        
    def build_index(self, data_dir="./data"):
        """构建索引"""
        try:
            # 加载文档
            documents = SimpleDirectoryReader(
                input_dir=data_dir,
                recursive=True
            ).load_data()
            
            logger.info(f"加载了 {len(documents)} 个文档")
            
            # 优化分块
            splitter = SentenceSplitter(
                chunk_size=512,
                chunk_overlap=50
            )
            nodes = splitter.get_nodes_from_documents(documents)
            
            # 构建索引
            self.index = VectorStoreIndex(
                nodes,
                vector_store=self.vector_store,
                show_progress=True
            )
            
            logger.info("索引构建完成")
            return True
            
        except Exception as e:
            logger.error(f"构建索引失败: {e}")
            return False
    
    def create_query_engine(self):
        """创建查询引擎"""
        # 检索器配置
        retriever = VectorIndexRetriever(
            index=self.index,
            similarity_top_k=20  # 粗排：取20个候选
        )
        
        # 后处理器（多层过滤）
        postprocessors = [
            SimilarityPostprocessor(similarity_cutoff=0.7),  # 过滤低相关
            CohereRerank(
                api_key="YOUR_KEY",
                top_n=5,  # 精排：选出5个
                model="rerank-multilingual-v3.0"
            )
        ]
        
        # 创建查询引擎
        query_engine = RetrieverQueryEngine(
            retriever=retriever,
            node_postprocessors=postprocessors,
            response_synthesizer=self.index.as_query_engine().response_synthesizer
        )
        
        return query_engine
    
    def query(self, question: str):
        """执行查询"""
        try:
            query_engine = self.create_query_engine()
            response = query_engine.query(question)
            
            # 返回结果
            return {
                "answer": str(response),
                "sources": [
                    {
                        "file": node.metadata.get("file_name"),
                        "score": node.score,
                        "text": node.text[:200]
                    }
                    for node in response.source_nodes
                ]
            }
        except Exception as e:
            logger.error(f"查询失败: {e}")
            return {"error": str(e)}

# 使用示例
if __name__ == "__main__":
    rag = ProductionRAG()
    
    # 构建索引（首次运行）
    rag.build_index("./data")
    
    # 查询
    result = rag.query("如何优化RAG检索质量？")
    print(f"回答: {result['answer']}")
    print(f"\n引用来源:")
    for src in result['sources']:
        print(f"  - {src['file']} (相关度: {src['score']:.3f})")