From c076023bc4b7c90f74b5b826f5dfdfb7d9c436f4 Mon Sep 17 00:00:00 2001
From: hailin <hailin@gdzx.xyz>
Date: Sat, 10 May 2025 12:58:46 +0800
Subject: [PATCH] .

---
 scripts/rag_build_query.py | 38 ++++++++++++++++++++++++++++++++------
 1 file changed, 32 insertions(+), 6 deletions(-)

diff --git a/scripts/rag_build_query.py b/scripts/rag_build_query.py
index 3df4562..958eb69 100644
--- a/scripts/rag_build_query.py
+++ b/scripts/rag_build_query.py
@@ -3,13 +3,14 @@ from typing import List  # 导入 List 用于类型注解
 import os
 import logging
 import faiss
-from llama_index import SimpleDirectoryReader, VectorStoreIndex, ServiceContext, StorageContext
+from llama_index import SimpleDirectoryReader, VectorStoreIndex, ServiceContext, StorageContext, load_index_from_storage
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.vector_stores.faiss import FaissVectorStore
 from app.core.config import settings  # 导入应用配置
 from scripts.permissions import get_user_allowed_indexes  # 导入权限函数，管理用户索引
 import numpy as np
 
+
 USER_INDEX_PATH = "index_data"  # 用户索引存储路径
 USER_DOC_PATH = "docs"  # 用户文档存储路径
 
@@ -152,22 +153,47 @@ def build_user_index(user_id: str):
 
 
 
+        # # 持久化存储之后，加载已保存的存储上下文信息
+        # logger.info(f"开始加载持久化存储的数据...")
+
+        # # 创建一个新的 StorageContext 实例，使用相同的目录
+        # loaded_storage_context = StorageContext.from_defaults(
+        #     persist_dir=persist_dir,  # 使用与之前相同的目录
+        #     vector_store=FaissVectorStore(faiss_index=faiss_index)  # 使用之前保存的 FAISS 索引
+        # )
+
+        # # 确认存储是否加载成功，检查索引数据
+        # logger.info("已成功加载存储上下文。")
+
+        # # 加载索引，进行检查
+        # loaded_index = VectorStoreIndex.from_storage_context(loaded_storage_context)
+        # logger.info(f"加载的索引数量: {len(loaded_index.get_documents())}")
+
+
         # 持久化存储之后，加载已保存的存储上下文信息
         logger.info(f"开始加载持久化存储的数据...")
 
+        # 加载 Faiss 索引
+        faiss_index_file = os.path.join(persist_dir, "index.faiss")
+        faiss_index = faiss.read_index(faiss_index_file)
+
+        # 创建 FaissVectorStore 实例
+        vector_store = FaissVectorStore(faiss_index=faiss_index)
+
         # 创建一个新的 StorageContext 实例，使用相同的目录
-        loaded_storage_context = StorageContext.from_defaults(
+        storage_context = StorageContext.from_defaults(
             persist_dir=persist_dir,  # 使用与之前相同的目录
-            vector_store=FaissVectorStore(faiss_index=faiss_index)  # 使用之前保存的 FAISS 索引
+            vector_store=vector_store  # 使用 Faiss 向量存储
         )
 
         # 确认存储是否加载成功，检查索引数据
         logger.info("已成功加载存储上下文。")
 
-        # 加载索引，进行检查
-        loaded_index = VectorStoreIndex.from_storage_context(loaded_storage_context)
-        logger.info(f"加载的索引数量: {len(loaded_index.get_documents())}")
+        # 使用 load_index_from_storage 加载索引
+        loaded_index = load_index_from_storage(storage_context)
 
+        # 检查加载后的索引
+        logger.info(f"加载的索引数量: {len(loaded_index.get_documents())}")
 
 
     except Exception as e: