LlamaIndex学习4-2：RAG索引存储

存储介绍

经过上一篇的测试，不管是长文档还是短文档，都会每次先进行向量化和建立索引，这里使用一下官方介绍的索引存储功能，来减少每次创建索引的步骤和模型消耗：

根据这篇文档说明：https://docs.llamaindex.ai/en/stable/module_guides/models/embeddings/#huggingface-optimum-onnx-embeddings

LlamaIndex 支持使用 HuggingFace 的 Optimum 库创建和使用 ONNX 嵌入。简单创建并保存 ONNX 嵌入，然后使用它们。

1. 使用HuggingFace的Optimum库创建ONNX嵌入模型
2. 实现索引的持久化存储和加载
3. 避免重复创建索引，提高系统效率

完整代码：

import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext, load_index_from_storage
from llama_index.llms.openai_like import OpenAILike
from llama_index.embeddings.huggingface_optimum import OptimumEmbedding
from llama_index.core import Settings
import asyncio
import shutil
from utils.logger import get_logger

# 获取封装好的日志工具
logger = get_logger()

# 定义目录和路径
ONNX_MODEL_DIR = "./onnx_model"
INDEX_STORAGE_DIR = "./storage"
DATA_DIR = "data"
HF_MODEL_NAME = "BAAI/bge-large-zh-v1.5"

def setup_llm():
    """设置LLM模型"""
    llm = OpenAILike(
        model="Qwen/Qwen2.5-7B-Instruct",
        api_key=os.getenv("SILICONFLOW_API_KEY"),
        api_base=os.getenv("SILICONFLOW_BASE_URL"),
        is_chat_model=True,
    )
    Settings.llm = llm
    logger.info("已设置全局LLM模型")
    return llm

def setup_optimum_embedding():
    """
    设置并获取OptimumEmbedding模型
    如果本地没有ONNX模型，则会创建并保存
    """
    # 检查是否已经有ONNX模型
    if not os.path.exists(ONNX_MODEL_DIR):
        logger.info(f"正在创建ONNX模型: {HF_MODEL_NAME} -> {ONNX_MODEL_DIR}")
        # 创建并保存ONNX模型
        OptimumEmbedding.create_and_save_optimum_model(
            HF_MODEL_NAME, 
            ONNX_MODEL_DIR
        )
        logger.info(f"ONNX模型已创建完成: {ONNX_MODEL_DIR}")
    else:
        logger.info(f"找到已有的ONNX模型: {ONNX_MODEL_DIR}")
    
    # 初始化OptimumEmbedding
    embed_model = OptimumEmbedding(folder_name=ONNX_MODEL_DIR)
    Settings.embed_model = embed_model
    logger.info("已设置全局Embedding模型")
    
    return embed_model

def get_or_create_index(embed_model, force_rebuild=False):
    """
    获取或创建向量索引
    Args:
        embed_model: 嵌入模型
        force_rebuild: 是否强制重建索引
    Returns:
        创建的索引
    """
    # 检查是否存在索引存储
    if os.path.exists(INDEX_STORAGE_DIR) and not force_rebuild:
        try:
            logger.info(f"正在从本地加载索引: {INDEX_STORAGE_DIR}")
            # 从存储加载索引
            storage_context = StorageContext.from_defaults(persist_dir=INDEX_STORAGE_DIR)
            index = load_index_from_storage(storage_context)
            logger.info("索引加载成功")
            return index
        except Exception as e:
            logger.error(f"加载索引失败: {e}")
            # 如果加载失败，删除可能已损坏的存储并重新创建
            logger.info("将重新创建索引")
            shutil.rmtree(INDEX_STORAGE_DIR, ignore_errors=True)
    
    # 创建新索引
    logger.info("正在创建新索引...")
    
    # 加载文档
    documents = SimpleDirectoryReader(DATA_DIR).load_data()
    logger.info(f"已加载 {len(documents)} 个文档")
    
    # 创建索引
    index = VectorStoreIndex.from_documents(
        documents,
        embed_model=embed_model,
        show_progress=True
    )
    
    # 持久化保存索引
    index.storage_context.persist(persist_dir=INDEX_STORAGE_DIR)
    logger.info(f"索引已创建并保存到: {INDEX_STORAGE_DIR}")
    
    return index

async def query_index(index, query_text):
    """使用索引查询文档"""
    logger.info(f"查询: {query_text}")
    
    # 创建查询引擎
    query_engine = index.as_query_engine()
    
    # 执行查询
    start_time = asyncio.get_event_loop().time()
    response = await query_engine.aquery(query_text)
    query_time = asyncio.get_event_loop().time() - start_time
    
    logger.info(f"查询完成，耗时: {query_time:.2f}秒")
    return response

async def main():
    """主函数"""
    # 设置LLM和嵌入模型
    setup_llm()
    embed_model = setup_optimum_embedding()
    
    # 获取或创建索引
    # 要强制重建索引，可以将force_rebuild设置为True
    index = get_or_create_index(embed_model, force_rebuild=False)
    
    # 执行测试查询
    test_query = "夏栖飞的真名叫什么？"
    response = await query_index(index, test_query)
    
    print("\n查询结果:")
    print(str(response))
    logger.info('查询结果：',str(response))

if __name__ == "__main__":
    logger.info("程序开始运行")
    asyncio.run(main())
    logger.info("程序运行结束")

注意事项

1. 首次运行：首次运行时会下载模型并创建ONNX优化版本，可能需要一些时间
2. 索引更新：当文档内容变更时，需要设置force_rebuild=True来重建索引
3. 存储路径：确保ONNX_MODEL_DIR和INDEX_STORAGE_DIR目录有写入权限
4. 异常处理：代码中包含了索引加载失败时的异常处理逻辑

查询效果展示

这里直接把《庆余年》这本原著小说进行了向量化和建立索引并存储起来，第一次访问会需要等待ONNX的进度，后续再重新执行py脚本，会直接从模型里进行查询：

问：夏栖飞的真名叫什么？
RAG回答

查询结果： 
夏栖飞的真名叫明青城。

问：范闲的箱子里藏着什么武器？
RAG回答

范闲的箱子里藏着一把名为M82A1的重狙。这是一把不属于这个世界、极其恐怖的武器，它曾导致庆国两位亲王的离奇死亡，帮助诚王爷登基，并让现在的庆国陛下有机会坐上龙椅。

什么是鱼肠？
RAG回答

查询结果： 鱼肠是一种特殊的剑，它被藏在鱼腹之中，可能永远不被发现，但如果被使用，就会刺入某个人的胸膛。范闲将自己和荆戈以及身边的影子比作鱼肠，表示他们虽然已经显露身份，但自己身边的鱼肠仍然隐藏未露。