import faiss import numpy as np import pickle import os from typing import List, Tuple, Optional, Union, Dict, Any import logging from pathlib import Path import json from datetime import datetime class FAISSManager: """ FAISS向量管理器 支持向量的增删改查、余弦相似度搜索、批量操作等功能 """ def __init__(self, dimension: int, index_type: str = "IndexFlatIP", normalize_vectors: bool = True, metric_type: str = "cosine"): """ 初始化FAISS管理器 Args: dimension: 向量维度 index_type: 索引类型 ("IndexFlatIP", "IndexFlatL2", "IndexIVFFlat", "IndexHNSW") normalize_vectors: 是否标准化向量（用于余弦相似度） metric_type: 距离度量类型 ("cosine", "l2", "ip") """ self.dimension = dimension self.index_type = index_type self.normalize_vectors = normalize_vectors self.metric_type = metric_type # 初始化索引 self.index = self._create_index() # 存储向量ID到实际ID的映射 self.id_mapping = {} # faiss_id -> actual_id self.reverse_id_mapping = {} # actual_id -> faiss_id self.next_faiss_id = 0 # 存储向量元数据 self.metadata = {} # actual_id -> metadata # 日志设置 logging.basicConfig(level=logging.INFO) self.logger = logging.getLogger(__name__) def _create_index(self) -> 'faiss.Index': """创建FAISS索引""" if self.index_type == "IndexFlatIP": # 内积索引（适合余弦相似度，需要标准化向量） # 初始化的时候需要告诉维度 index = faiss.IndexFlatIP(self.dimension) elif self.index_type == "IndexFlatL2": # L2距离索引 index = faiss.IndexFlatL2(self.dimension) elif self.index_type == "IndexIVFFlat": # IVF索引（适合大规模数据） quantizer = faiss.IndexFlatIP(self.dimension) if self.metric_type == "cosine" else faiss.IndexFlatL2(self.dimension) nlist = 100 # 聚类中心数量 index = faiss.IndexIVFFlat(quantizer, self.dimension, nlist) elif self.index_type == "IndexHNSW": # HNSW索引（高性能近似搜索） index = faiss.IndexHNSWFlat(self.dimension, 32) else: raise ValueError(f"不支持的索引类型: {self.index_type}") return index def _normalize_vector(self, vector: np.ndarray) -> np.ndarray: """标准化向量（用于余弦相似度）""" if self.normalize_vectors: norm = np.linalg.norm(vector, axis=-1, keepdims=True) # 避免除零 norm = np.where(norm == 0, 1, norm) return vector / norm return vector def add_vector(self, vector: Union[np.ndarray, List[float]], vector_id: str, metadata: Optional[Dict[str, Any]] = None) -> bool: """ 添加单个向量 Args: vector: 向量数据 vector_id: 向量唯一标识 metadata: 向量元数据 Returns: bool: 是否添加成功 """ try: # 检查向量是否已存在 if vector_id in self.reverse_id_mapping: self.logger.warning(f"向量ID {vector_id} 已存在，将更新该向量") return self.update_vector(vector, vector_id, metadata) # 转换为numpy数组并标准化 vector = np.array(vector, dtype=np.float32).reshape(1, -1) if vector.shape[1] != self.dimension: raise ValueError(f"向量维度不匹配: 期望 {self.dimension}, 实际 {vector.shape[1]}") vector = self._normalize_vector(vector) # 添加到索引 self.index.add(vector) # 更新映射关系 faiss_id = self.next_faiss_id self.id_mapping[faiss_id] = vector_id self.reverse_id_mapping[vector_id] = faiss_id self.next_faiss_id += 1 # 存储元数据 if metadata: self.metadata[vector_id] = metadata self.logger.info(f"成功添加向量: {vector_id}") return True except Exception as e: self.logger.error(f"添加向量失败: {e}") return False def add_vectors_batch(self, vectors: Union[np.ndarray, List[List[float]]], vector_ids: List[str], metadata_list: Optional[List[Dict[str, Any]]] = None) -> List[bool]: """ 批量添加向量 Args: vectors: 向量数据矩阵 vector_ids: 向量ID列表 metadata_list: 元数据列表 Returns: List[bool]: 每个向量的添加结果 """ results = [] # 转换为numpy数组 vectors = np.array(vectors, dtype=np.float32) if vectors.ndim == 1: vectors = vectors.reshape(1, -1) if vectors.shape[1] != self.dimension: raise ValueError(f"向量维度不匹配: 期望 {self.dimension}, 实际 {vectors.shape[1]}") if len(vector_ids) != vectors.shape[0]: raise ValueError("向量数量与ID数量不匹配") # 标准化向量 vectors = self._normalize_vector(vectors) # 检查重复ID new_vectors = [] new_ids = [] new_metadata = [] for i, vector_id in enumerate(vector_ids): if vector_id not in self.reverse_id_mapping: new_vectors.append(vectors[i]) new_ids.append(vector_id) if metadata_list: new_metadata.append(metadata_list[i]) results.append(True) else: self.logger.warning(f"向量ID {vector_id} 已存在，跳过") results.append(False) if new_vectors: # 批量添加到索引 new_vectors = np.array(new_vectors) self.index.add(new_vectors) # 批量更新映射关系 for i, vector_id in enumerate(new_ids): faiss_id = self.next_faiss_id + i self.id_mapping[faiss_id] = vector_id self.reverse_id_mapping[vector_id] = faiss_id # 存储元数据 if new_metadata and i < len(new_metadata): self.metadata[vector_id] = new_metadata[i] self.next_faiss_id += len(new_ids) self.logger.info(f"批量添加了 {len(new_ids)} 个向量") return results def search_similar(self, query_vector: Union[np.ndarray, List[float]], k: int = 10, threshold: Optional[float] = None) -> List[Tuple[str, float, Dict[str, Any]]]: """ 搜索相似向量 Args: query_vector: 查询向量 k: 返回最相似的k个结果 threshold: 相似度阈值 Returns: List[Tuple[str, float, Dict]]: (vector_id, similarity_score, metadata) """ if self.index.ntotal == 0: return [] # 转换并标准化查询向量 query_vector = np.array(query_vector, dtype=np.float32).reshape(1, -1) if query_vector.shape[1] != self.dimension: raise ValueError(f"查询向量维度不匹配: 期望 {self.dimension}, 实际 {query_vector.shape[1]}") query_vector = self._normalize_vector(query_vector) # 搜索 k = min(k, self.index.ntotal) distances, indices = self.index.search(query_vector, k) results = [] for i in range(k): faiss_id = indices[0][i] distance = distances[0][i] if faiss_id == -1: # 无效结果 continue # 转换距离为相似度分数 if self.metric_type == "cosine" or self.index_type == "IndexFlatIP": similarity = float(distance) # 内积已经是相似度 else: similarity = 1.0 / (1.0 + float(distance)) # L2距离转相似度 # 应用阈值过滤 if threshold is not None and similarity < threshold: continue vector_id = self.id_mapping.get(faiss_id) if vector_id: metadata = self.metadata.get(vector_id, {}) results.append((vector_id, similarity, metadata)) return results def get_vector_by_id(self, vector_id: str) -> Optional[Tuple[np.ndarray, Dict[str, Any]]]: """ 根据ID获取向量 Args: vector_id: 向量ID Returns: Optional[Tuple[np.ndarray, Dict]]: (vector, metadata) 或 None """ faiss_id = self.reverse_id_mapping.get(vector_id) if faiss_id is None: return None # FAISS不直接支持根据ID获取向量，这里返回元数据 # 如果需要向量数据，建议单独存储 metadata = self.metadata.get(vector_id, {}) return None, metadata def update_vector(self, vector: Union[np.ndarray, List[float]], vector_id: str, metadata: Optional[Dict[str, Any]] = None) -> bool: """ 更新向量（先删除再添加） Args: vector: 新向量数据 vector_id: 向量ID metadata: 新元数据 Returns: bool: 是否更新成功 """ # FAISS不支持直接更新，需要重建索引 if vector_id not in self.reverse_id_mapping: self.logger.warning(f"向量ID {vector_id} 不存在") return False # 删除旧向量 self.delete_vector(vector_id) # 添加新向量 return self.add_vector(vector, vector_id, metadata) def delete_vector(self, vector_id: str) -> bool: """ 删除向量 Args: vector_id: 向量ID Returns: bool: 是否删除成功 """ if vector_id not in self.reverse_id_mapping: self.logger.warning(f"向量ID {vector_id} 不存在") return False # FAISS不支持直接删除，标记为删除 faiss_id = self.reverse_id_mapping[vector_id] # 从映射中移除 del self.id_mapping[faiss_id] del self.reverse_id_mapping[vector_id] # 删除元数据 if vector_id in self.metadata: del self.metadata[vector_id] self.logger.info(f"标记删除向量: {vector_id}") return True def delete_vectors_batch(self, vector_ids: List[str]) -> List[bool]: """ 批量删除向量 Args: vector_ids: 向量ID列表 Returns: List[bool]: 每个向量的删除结果 """ results = [] for vector_id in vector_ids: results.append(self.delete_vector(vector_id)) return results def rebuild_index(self) -> bool: """ 重建索引（清理已删除的向量） """ try: # 获取所有有效向量的信息 valid_vectors = [] valid_ids = [] valid_metadata = {} # 这里需要外部提供向量数据，因为FAISS不能直接获取 # 建议在实际使用时维护一个向量存储 self.logger.warning("重建索引需要外部提供向量数据") return False except Exception as e: self.logger.error(f"重建索引失败: {e}") return False def get_index_info(self) -> Dict[str, Any]: """ 获取索引信息 Returns: Dict: 索引统计信息 """ return { "total_vectors": self.index.ntotal, "dimension": self.dimension, "index_type": self.index_type, "metric_type": self.metric_type, "normalize_vectors": self.normalize_vectors, "active_vectors": len(self.reverse_id_mapping), "metadata_count": len(self.metadata) } def save_index(self, save_path: str) -> bool: """ 保存索引到文件 Args: save_path: 保存路径 Returns: bool: 是否保存成功 """ try: save_path = Path(save_path) save_path.mkdir(parents=True, exist_ok=True) # 保存FAISS索引 index_file = save_path / "faiss_index.bin" faiss.write_index(self.index, str(index_file)) # 保存映射关系和元数据 metadata_file = save_path / "metadata.pkl" with open(metadata_file, 'wb') as f: pickle.dump({ 'id_mapping': self.id_mapping, 'reverse_id_mapping': self.reverse_id_mapping, 'next_faiss_id': self.next_faiss_id, 'metadata': self.metadata, 'dimension': self.dimension, 'index_type': self.index_type, 'normalize_vectors': self.normalize_vectors, 'metric_type': self.metric_type }, f) # 保存配置信息 config_file = save_path / "config.json" with open(config_file, 'w', encoding='utf-8') as f: json.dump({ 'dimension': self.dimension, 'index_type': self.index_type, 'normalize_vectors': self.normalize_vectors, 'metric_type': self.metric_type, 'save_time': datetime.now().isoformat(), 'total_vectors': self.index.ntotal }, f, indent=2, ensure_ascii=False) self.logger.info(f"索引已保存到: {save_path}") return True except Exception as e: self.logger.error(f"保存索引失败: {e}") return False def load_index(self, load_path: str) -> bool: """ 从文件加载索引 Args: load_path: 加载路径 Returns: bool: 是否加载成功 """ try: load_path = Path(load_path) # 检查文件是否存在 index_file = load_path / "faiss_index.bin" metadata_file = load_path / "metadata.pkl" if not index_file.exists() or not metadata_file.exists(): self.logger.error(f"索引文件不存在: {load_path}") return False # 加载FAISS索引 self.index = faiss.read_index(str(index_file)) # 加载映射关系和元数据 with open(metadata_file, 'rb') as f: data = pickle.load(f) self.id_mapping = data['id_mapping'] self.reverse_id_mapping = data['reverse_id_mapping'] self.next_faiss_id = data['next_faiss_id'] self.metadata = data['metadata'] self.dimension = data['dimension'] self.index_type = data['index_type'] self.normalize_vectors = data['normalize_vectors'] self.metric_type = data['metric_type'] self.logger.info(f"索引已从 {load_path} 加载") return True except Exception as e: self.logger.error(f"加载索引失败: {e}") return False def clear_index(self) -> bool: """ 清空索引 Returns: bool: 是否清空成功 """ try: # 重新创建索引 self.index = self._create_index() # 清空映射关系 self.id_mapping.clear() self.reverse_id_mapping.clear() self.next_faiss_id = 0 # 清空元数据 self.metadata.clear() self.logger.info("索引已清空") return True except Exception as e: self.logger.error(f"清空索引失败: {e}") return False # 使用示例 if __name__ == "__main__": # 创建FAISS管理器 manager = FAISSManager(dimension=512, index_type="IndexFlatIP") # 添加向量 vector1 = np.random.random(512).astype(np.float32) manager.add_vector(vector1, "vec_1", {"label": "食物1", "category": "川菜"}) # 批量添加向量 vectors = np.random.random((10, 512)).astype(np.float32) ids = [f"vec_{i}" for i in range(2, 12)] metadata_list = [{"label": f"食物{i}", "category": "川菜"} for i in range(2, 12)] manager.add_vectors_batch(vectors, ids, metadata_list) # 搜索相似向量 query = np.random.random(512).astype(np.float32) results = manager.search_similar(query, k=5) print("搜索结果:", results) # 获取索引信息 # 获得统计信息 info = manager.get_index_info() print("索引信息:", info) # 保存索引 manager.save_index("./faiss_index") # 加载索引 new_manager = FAISSManager(dimension=512) new_manager.load_index("./faiss_index")