Initial commit

2026-05-12 09:41:56 +08:00
commit 572283e101
936 changed files with 133949 additions and 0 deletions
@@ -0,0 +1,135 @@
+"""
+索引构建模块
+"""
+
+import logging
+from typing import List
+from pathlib import Path
+
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_core.documents import Document
+
+logger = logging.getLogger(__name__)
+
+class IndexConstructionModule:
+    """索引构建模块 - 负责向量化和索引构建"""
+
+    def __init__(self, model_name: str = "BAAI/bge-small-zh-v1.5", index_save_path: str = "./vector_index"):
+        """
+        初始化索引构建模块
+
+        Args:
+            model_name: 嵌入模型名称
+            index_save_path: 索引保存路径
+        """
+        self.model_name = model_name
+        self.index_save_path = index_save_path
+        self.embeddings = None
+        self.vectorstore = None
+        self.setup_embeddings()
+    
+    def setup_embeddings(self):
+        """初始化嵌入模型"""
+        logger.info(f"正在初始化嵌入模型: {self.model_name}")
+        
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name=self.model_name,
+            model_kwargs={'device': 'cpu'},
+            encode_kwargs={'normalize_embeddings': True}
+        )
+        
+        logger.info("嵌入模型初始化完成")
+    
+    def build_vector_index(self, chunks: List[Document]) -> FAISS:
+        """
+        构建向量索引
+        
+        Args:
+            chunks: 文档块列表
+            
+        Returns:
+            FAISS向量存储对象
+        """
+        logger.info("正在构建FAISS向量索引...")
+        
+        if not chunks:
+            raise ValueError("文档块列表不能为空")
+        
+        # 构建FAISS向量存储
+        self.vectorstore = FAISS.from_documents(
+            documents=chunks,
+            embedding=self.embeddings
+        )
+        
+        logger.info(f"向量索引构建完成，包含 {len(chunks)} 个向量")
+        return self.vectorstore
+    
+    def add_documents(self, new_chunks: List[Document]):
+        """
+        向现有索引添加新文档
+        
+        Args:
+            new_chunks: 新的文档块列表
+        """
+        if not self.vectorstore:
+            raise ValueError("请先构建向量索引")
+        
+        logger.info(f"正在添加 {len(new_chunks)} 个新文档到索引...")
+        self.vectorstore.add_documents(new_chunks)
+        logger.info("新文档添加完成")
+
+    def save_index(self):
+        """
+        保存向量索引到配置的路径
+        """
+        if not self.vectorstore:
+            raise ValueError("请先构建向量索引")
+
+        # 确保保存目录存在
+        Path(self.index_save_path).mkdir(parents=True, exist_ok=True)
+
+        self.vectorstore.save_local(self.index_save_path)
+        logger.info(f"向量索引已保存到: {self.index_save_path}")
+    
+    def load_index(self):
+        """
+        从配置的路径加载向量索引
+
+        Returns:
+            加载的向量存储对象，如果加载失败返回None
+        """
+        if not self.embeddings:
+            self.setup_embeddings()
+
+        if not Path(self.index_save_path).exists():
+            logger.info(f"索引路径不存在: {self.index_save_path}，将构建新索引")
+            return None
+
+        try:
+            self.vectorstore = FAISS.load_local(
+                self.index_save_path,
+                self.embeddings,
+                allow_dangerous_deserialization=True
+            )
+            logger.info(f"向量索引已从 {self.index_save_path} 加载")
+            return self.vectorstore
+        except Exception as e:
+            logger.warning(f"加载向量索引失败: {e}，将构建新索引")
+            return None
+    
+    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
+        """
+        相似度搜索
+        
+        Args:
+            query: 查询文本
+            k: 返回结果数量
+            
+        Returns:
+            相似文档列表
+        """
+        if not self.vectorstore:
+            raise ValueError("请先构建或加载向量索引")
+        
+        return self.vectorstore.similarity_search(query, k=k)