This commit is contained in:
hailin 2025-05-10 12:58:46 +08:00
parent 19e00f1728
commit c076023bc4
1 changed files with 32 additions and 6 deletions

View File

@ -3,13 +3,14 @@ from typing import List # 导入 List 用于类型注解
import os import os
import logging import logging
import faiss 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.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.vector_stores.faiss import FaissVectorStore from llama_index.vector_stores.faiss import FaissVectorStore
from app.core.config import settings # 导入应用配置 from app.core.config import settings # 导入应用配置
from scripts.permissions import get_user_allowed_indexes # 导入权限函数,管理用户索引 from scripts.permissions import get_user_allowed_indexes # 导入权限函数,管理用户索引
import numpy as np import numpy as np
USER_INDEX_PATH = "index_data" # 用户索引存储路径 USER_INDEX_PATH = "index_data" # 用户索引存储路径
USER_DOC_PATH = "docs" # 用户文档存储路径 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"开始加载持久化存储的数据...") 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 实例,使用相同的目录 # 创建一个新的 StorageContext 实例,使用相同的目录
loaded_storage_context = StorageContext.from_defaults( storage_context = StorageContext.from_defaults(
persist_dir=persist_dir, # 使用与之前相同的目录 persist_dir=persist_dir, # 使用与之前相同的目录
vector_store=FaissVectorStore(faiss_index=faiss_index) # 使用之前保存的 FAISS 索引 vector_store=vector_store # 使用 Faiss 向量存储
) )
# 确认存储是否加载成功,检查索引数据 # 确认存储是否加载成功,检查索引数据
logger.info("已成功加载存储上下文。") logger.info("已成功加载存储上下文。")
# 加载索引,进行检查 # 使用 load_index_from_storage 加载索引
loaded_index = VectorStoreIndex.from_storage_context(loaded_storage_context) loaded_index = load_index_from_storage(storage_context)
logger.info(f"加载的索引数量: {len(loaded_index.get_documents())}")
# 检查加载后的索引
logger.info(f"加载的索引数量: {len(loaded_index.get_documents())}")
except Exception as e: except Exception as e: