isage-middleware 0.1.4__py3-none-any.whl → 0.1.5.1__py3-none-any.whl

sage/middleware/services/vdb/search_engine/faiss_index.py

@@ -1,461 +0,0 @@
-# file sage/core/sage.service.memory./search_engine/vdb_index/faiss_index.py
-# python -m sage.core.sage.service.memory..search_engine.vdb_index.faiss_index
-
-import json
-import faiss
-import pickle
-import numpy as np
-from typing import Optional, List, Dict, Any
-from sage.common.utils.logging.custom_logger import CustomLogger
-from sage.middleware.services.vdb.search_engine.base_vdb_index import BaseVDBIndex
-import os
-
-class FaissIndex(BaseVDBIndex):
-    def __init__(
-        self, 
-        name: str, 
-        dim: int, 
-        vectors: Optional[List[np.ndarray]] = None, 
-        ids: Optional[List[str]] = None,
-        config: Optional[dict] = None,
-        load_path: Optional[str] = None,
-        ):
-        """
-        初始化 FaissIndex 实例，设置索引名称、维度，并可选性地加载初始向量和ID
-        Initialize the FaissIndex instance with name, dimension, and optionally preload vectors and ids.
-        """
-        self.index_name = name
-        self.dim = dim
-        self.config = config or {}
-        self.id_map: Dict[int, str] = {}
-        self.rev_map: Dict[str, int] = {}
-        self.next_id: int = 1
-        self.tombstones: set[str] = set()
-        self._deletion_supported = True
-        self.index = None
-        self.logger = CustomLogger()
-        if load_path is not None:
-            self._load(load_path)
-        else:
-            self.index, self._deletion_supported = self._init_index()
-            if vectors is not None and ids is not None:
-                self._build_index(vectors, ids)
-
-           
-    def _init_index(self):
-        config = self.config  # 保持全程都叫config
-        index_type = config.get("index_type", "IndexFlatL2")
-
-        # 基础索引
-        if index_type == "IndexFlatL2":
-            return faiss.IndexFlatL2(self.dim), True
-
-        elif index_type == "IndexFlatIP":
-            return faiss.IndexFlatIP(self.dim), True
-
-        # HNSW
-        elif index_type == "IndexHNSWFlat":
-            hnsw_m = int(config.get("HNSW_M", 32))
-            ef_construction = int(config.get("HNSW_EF_CONSTRUCTION", 200))
-            index = faiss.IndexHNSWFlat(self.dim, hnsw_m)
-            index.hnsw.efConstruction = ef_construction
-            if "HNSW_EF_SEARCH" in config:
-                index.hnsw.efSearch = int(config["HNSW_EF_SEARCH"])
-            return index, False
-
-        # IVF Flat
-        elif index_type == "IndexIVFFlat":
-            nlist = int(config.get("IVF_NLIST", 100))
-            nprobe = int(config.get("IVF_NPROBE", 10))
-            metric = self._get_metric(config.get("IVF_METRIC", "L2"))
-            quantizer = faiss.IndexFlatL2(self.dim)
-            index = faiss.IndexIVFFlat(quantizer, self.dim, nlist, metric)
-            index.nprobe = nprobe
-            return index, True
-
-        # IVF PQ
-        elif index_type == "IndexIVFPQ":
-            nlist = int(config.get("IVF_NLIST", 100))
-            nprobe = int(config.get("IVF_NPROBE", 10))
-            m = int(config.get("PQ_M", 8))
-            nbits = int(config.get("PQ_NBITS", 8))
-            metric = self._get_metric(config.get("IVF_METRIC", "L2"))
-            quantizer = faiss.IndexFlatL2(self.dim)
-            index = faiss.IndexIVFPQ(quantizer, self.dim, nlist, m, nbits, metric)
-            index.nprobe = nprobe
-            return index, True
-
-        # IVF ScalarQuantizer
-        elif index_type == "IndexIVFScalarQuantizer":
-            nlist = int(config.get("IVF_NLIST", 100))
-            nprobe = int(config.get("IVF_NPROBE", 10))
-            qtype_str = config.get("SQ_TYPE", "QT_8bit")
-            qtype = getattr(faiss.ScalarQuantizer, qtype_str)
-            metric = self._get_metric(config.get("IVF_METRIC", "L2"))
-            quantizer = faiss.IndexFlatL2(self.dim)
-            index = faiss.IndexIVFScalarQuantizer(quantizer, self.dim, nlist, qtype, metric)
-            index.nprobe = nprobe
-            return index, True
-
-        # LSH
-        elif index_type == "IndexLSH":
-            nbits = int(config.get("LSH_NBITS", 512))
-            rotate_data = bool(config.get("LSH_ROTATE_DATA", True))
-            train_thresholds = bool(config.get("LSH_TRAIN_THRESHOLDS", False))
-            index = faiss.IndexLSH(self.dim, nbits, rotate_data, train_thresholds)
-            return index, False
-
-        # PQ
-        elif index_type == "IndexPQ":
-            m = int(config.get("PQ_M", 8))
-            nbits = int(config.get("PQ_NBITS", 8))
-            metric = self._get_metric(config.get("PQ_METRIC", "L2"))
-            return faiss.IndexPQ(self.dim, m, nbits, metric), False
-
-        # ScalarQuantizer
-        elif index_type == "IndexScalarQuantizer":
-            qtype_str = config.get("SQ_TYPE", "QT_8bit")
-            qtype = getattr(faiss.ScalarQuantizer, qtype_str)
-            metric = self._get_metric(config.get("SQ_METRIC", "L2"))
-            return faiss.IndexScalarQuantizer(self.dim, qtype, metric), True
-
-        # RefineFlat
-        elif index_type == "IndexRefineFlat":
-            base_type = config.get("FAISS_BASE_INDEX_TYPE", "IndexFlatL2")
-            # 临时切换 index_type, 递归用 config 初始化
-            orig_type = config.get("index_type", None)
-            config["index_type"] = base_type
-            base_index, base_deletion_supported = self._init_index()
-            if orig_type is not None:
-                config["index_type"] = orig_type
-            k_factor = float(config.get("REFINE_K_FACTOR", 1.0))
-            return faiss.IndexRefineFlat(base_index, k_factor), True
-
-        # IndexIDMap
-        elif index_type == "IndexIDMap":
-            base_type = config.get("FAISS_BASE_INDEX_TYPE", "IndexFlatL2")
-            orig_type = config.get("index_type", None)
-            config["index_type"] = base_type
-            base_index, base_deletion_supported = self._init_index()
-            if orig_type is not None:
-                config["index_type"] = orig_type
-            return faiss.IndexIDMap(base_index), base_deletion_supported
-
-        else:
-            raise ValueError(f"Unsupported FAISS index type: {index_type}")
-
-    def _init_base_index(self):
-        """
-        用于 IndexIDMap / IndexRefineFlat 的基础索引初始化
-        Initialize base index for IndexIDMap or IndexRefineFlat
-        """
-        base_type = os.getenv("FAISS_BASE_INDEX_TYPE", "IndexFlatL2")
-
-        original_type = os.getenv("FAISS_INDEX_TYPE")
-        os.environ["FAISS_INDEX_TYPE"] = base_type
-        index = self._init_index()
-        if original_type:
-            os.environ["FAISS_INDEX_TYPE"] = original_type
-        return index
-
-    def _get_metric(self, metric_str):
-        """
-        获取距离度量方式：L2 或 Inner Product
-        Get distance metric: L2 or Inner Product
-        """
-        return faiss.METRIC_L2 if metric_str == "L2" else faiss.METRIC_INNER_PRODUCT
-    
-    def _build_index(self, vectors: List[np.ndarray], ids: List[str]):
-        """
-        构建初始索引并绑定 string ID → int ID 映射关系
-        Build initial index and bind string ID to int ID mapping
-        """
-        np_vectors = np.vstack(vectors).astype("float32")
-        int_ids = []
-
-        for string_id in ids:
-            if string_id in self.rev_map:
-                int_id = self.rev_map[string_id]
-            else:
-                int_id = self.next_id
-                self.next_id += 1
-                self.rev_map[string_id] = int_id
-                self.id_map[int_id] = string_id
-            int_ids.append(int_id)
-
-        int_ids_np = np.array(int_ids, dtype=np.int64)
-        if not isinstance(self.index, faiss.IndexIDMap):
-            self.logger.info("Wrapping index with IndexIDMap")
-            self.index = faiss.IndexIDMap(self.index)  # 仅当未包装时才包装
-        self.index.add_with_ids(np_vectors, int_ids_np)  # type: ignore
-        
-    def delete(self, string_id: str):
-        """
-        删除指定ID（物理删除或墓碑标记）
-        Delete by ID (physical removal or tombstone marking)
-        """
-        if string_id not in self.rev_map:
-            return
-
-        int_id = self.rev_map[string_id]
-        
-        if self._deletion_supported:
-            try:
-                id_vector = np.array([int_id], dtype=np.int64)
-                self.index.remove_ids(id_vector)  # type: ignore
-            except Exception as e:
-                print(f"删除失败，转为墓碑标记 / Deletion failed, fallback to tombstone: {e}")
-                self.tombstones.add(string_id)
-        else:
-            self.tombstones.add(string_id)
-    
-    def update(self, string_id: str, new_vector: np.ndarray):
-        """
-        更新指定 ID 的向量：保持原有映射关系，仅替换向量内容
-        Update the vector for the given ID, preserving the existing ID mapping.
-        """
-        if string_id not in self.rev_map:
-            # 如果ID不存在，直接插入
-            return self.insert(new_vector, string_id)
-
-        int_id = self.rev_map[string_id]
-
-        if self._deletion_supported:
-            try:
-                 # 删除旧向量并插入新向量 / Remove old vector and insert new one
-                id_vector = np.array([int_id], dtype=np.int64)
-                self.index.remove_ids(id_vector)  # type: ignore
-                vector = np.expand_dims(new_vector.astype("float32"), axis=0)
-                int_id_np = np.array([int_id], dtype=np.int64)
-                self.index.add_with_ids(vector, int_id_np)  # type: ignore
-            except Exception as e:
-                print(f"更新失败 / Update failed: {e}")
-                self.insert(new_vector, string_id)
-        else:
-            if string_id in self.rev_map:
-                old_int_id = self.rev_map[string_id]
-                if old_int_id in self.id_map:
-                    del self.id_map[old_int_id]
-                del self.rev_map[string_id]
-            
-            new_int_id = self.next_id
-            self.next_id += 1
-            self.rev_map[string_id] = new_int_id
-            self.id_map[new_int_id] = string_id
-            vector = np.expand_dims(new_vector.astype("float32"), axis=0)
-            int_id_np = np.array([new_int_id], dtype=np.int64)
-            self.index.add_with_ids(vector, int_id_np)  # type: ignore
-
-    def search(self, query_vector: np.ndarray, topk: int = 10):
-        """
-        向量检索 / Vector search
-        返回top_k结果（过滤墓碑） / Return top_k results (filter tombstones)
-        """
-        query_vector = np.expand_dims(query_vector.astype("float32"), axis=0)
-        distances, int_ids = self.index.search(query_vector, topk + len(self.tombstones))  # 多查一些结果 # type: ignore
-        
-        results = []
-        filtered_distances = []
-        
-        for i, dist in zip(int_ids[0], distances[0]):
-            if i == -1:  # FAISS 空槽位标记
-                continue
-            string_id = self.id_map.get(i)
-            if string_id and string_id not in self.tombstones:
-                results.append(string_id)
-                filtered_distances.append(float(dist))  # 显式转为Python float
-            if len(results) >= topk:
-                break
-                
-        return results, filtered_distances
-
-    def insert(self, vector: np.ndarray, string_id: str):
-        """
-        插入单个向量及其字符串 ID 到索引中
-        Insert a single vector and its string ID into the index
-        """
-        if string_id in self.rev_map:
-            int_id = self.rev_map[string_id]
-        else:
-            int_id = self.next_id
-            self.next_id += 1
-            self.rev_map[string_id] = int_id
-            self.id_map[int_id] = string_id
-
-        vector = np.expand_dims(vector.astype("float32"), axis=0)
-        int_id_np = np.array([int_id], dtype=np.int64)
-        self.index.add_with_ids(vector, int_id_np) # type: ignore
-    
-    def batch_insert(self, vectors: List[np.ndarray], string_ids: List[str]):
-        """
-        批量插入多个向量及其对应的 string_id
-        Batch insert multiple vectors and their corresponding string_id
-        """
-        assert len(vectors) == len(string_ids), "Vectors and IDs must match in length"
-        np_vectors = np.vstack(vectors).astype("float32")
-        int_ids = []
-
-        for string_id in string_ids:
-            if string_id in self.rev_map:
-                int_id = self.rev_map[string_id]
-            else:
-                int_id = self.next_id
-                self.next_id += 1
-                self.rev_map[string_id] = int_id
-                self.id_map[int_id] = string_id
-            int_ids.append(int_id)
-
-        int_ids_np = np.array(int_ids, dtype=np.int64)
-        self.index.add_with_ids(np_vectors, int_ids_np)  # type: ignore
-
-    def store(self, dir_path: str) -> Dict[str, Any]:
-            """
-            将FAISS索引、参数和映射全部保存到指定目录。
-            """
-            os.makedirs(dir_path, exist_ok=True)
-            # 1. 保存faiss主索引
-            faiss.write_index(self.index, os.path.join(dir_path, "faiss.index"))
-            # 2. 保存id映射
-            with open(os.path.join(dir_path, "id_map.pkl"), "wb") as f:
-                pickle.dump(self.id_map, f)
-            with open(os.path.join(dir_path, "rev_map.pkl"), "wb") as f:
-                pickle.dump(self.rev_map, f)
-            with open(os.path.join(dir_path, "tombstones.pkl"), "wb") as f:
-                pickle.dump(self.tombstones, f)
-            # 3. 保存参数（如dim、下一个ID、自定义config等）
-            meta = {
-                "index_name": self.index_name,
-                "dim": self.dim,
-                "next_id": self.next_id,
-                "deletion_supported": self._deletion_supported,
-                "config": getattr(self, "config", {}),  # 若有config则保存
-            }
-            with open(os.path.join(dir_path, "meta.json"), "w", encoding="utf-8") as f:
-                json.dump(meta, f, ensure_ascii=False, indent=2)
-            return {"index_path": dir_path}
-
-    def _load(self, dir_path: str):
-        """
-        从目录恢复索引、映射与参数。仅供类方法load调用。
-        """
-        self.index = faiss.read_index(os.path.join(dir_path, "faiss.index"))
-        with open(os.path.join(dir_path, "id_map.pkl"), "rb") as f:
-            self.id_map = pickle.load(f)
-        with open(os.path.join(dir_path, "rev_map.pkl"), "rb") as f:
-            self.rev_map = pickle.load(f)
-        with open(os.path.join(dir_path, "tombstones.pkl"), "rb") as f:
-            self.tombstones = pickle.load(f)
-        with open(os.path.join(dir_path, "meta.json"), "r", encoding="utf-8") as f:
-            meta = json.load(f)
-        self.index_name = meta["index_name"]
-        self.dim = meta["dim"]
-        self.next_id = meta["next_id"]
-        self._deletion_supported = meta.get("deletion_supported", True)
-        self.config = meta.get("config", {})
-
-    @classmethod
-    def load(cls, name: str, load_path: str) -> "FaissIndex":
-        """
-        推荐方式，等价于 BM25sIndex.load
-        """
-        return cls(name=name, dim=0, load_path=load_path)  # dim会被_load覆盖
-
-if __name__ == "__main__":
-    import os
-    import shutil
-    import numpy as np
-
-    def colored(text, color):
-        # color: "green", "red", "yellow"
-        colors = {"green": "\033[92m", "red": "\033[91m", "yellow": "\033[93m", "reset": "\033[0m"}
-        return colors.get(color, "") + text + colors["reset"]
-
-    def print_test_case(desc, expected_ids, expected_dists, actual_ids, actual_dists, digits=4):
-        ids_pass = list(expected_ids) == list(actual_ids)
-        dists_pass = all(abs(e-a) < 10**-digits for e,a in zip(expected_dists, actual_dists))
-        status = "通过" if ids_pass and dists_pass else "不通过"
-        color = "green" if status == "通过" else "red"
-        print(f"【{desc}】")
-        print(f"预期IDs：{expected_ids}")
-        print(f"实际IDs：{actual_ids}")
-        print(f"预期距离：{expected_dists}")
-        print(f"实际距离：{[round(x, digits) for x in actual_dists]}")
-        print(f"测试情况：{colored(status, color)}\n")
-
-    # ==== 基础数据 ====
-    dim = 4
-    index_name = "test_index"
-    root_dir = "./faiss_index_test"
-    if os.path.exists(root_dir):
-        shutil.rmtree(root_dir)
-    os.makedirs(root_dir, exist_ok=True)
-
-    vectors = [
-        np.array([1.0, 0.0, 0.0, 0.0]),
-        np.array([0.0, 1.0, 0.0, 0.0]),
-        np.array([0.0, 0.0, 1.0, 0.0])
-    ]
-    ids = ["id1", "id2", "id3"]
-
-    index = FaissIndex(name=index_name, dim=dim, vectors=vectors, ids=ids)
-    # 1. 检索
-    q1 = np.array([1.0, 0.0, 0.0, 0.0])
-    r_ids, r_dists = index.search(q1, 3)
-    print_test_case("基础检索", ["id1", "id2", "id3"], [0.0, 2.0, 2.0], r_ids, r_dists)
-
-    # 2. 插入新向量
-    index.insert(np.array([0.0, 0.0, 0.0, 1.0]), "id4")
-    q2 = np.array([0.0, 0.0, 0.0, 1.0])
-    r_ids, r_dists = index.search(q2, 4)
-    print_test_case("插入后检索", ["id4", "id1", "id2", "id3"], [0.0, 2.0, 2.0, 2.0], r_ids, r_dists)
-
-    # 3. 更新向量
-    index.update("id1", np.array([0.5, 0.5, 0.0, 0.0]))
-    q3 = np.array([0.5, 0.5, 0.0, 0.0])
-    r_ids, r_dists = index.search(q3, 4)
-    print_test_case("更新后检索", ['id1', 'id2', 'id3', 'id4'], [0.0, 0.5, 1.5, 1.5], r_ids, r_dists)
-
-    # 4. 删除向量
-    index.delete("id2")
-    q4 = np.array([1.0, 0.0, 0.0, 0.0])
-    r_ids, r_dists = index.search(q4, 4)
-    print_test_case("删除后检索", ['id1', 'id3', 'id4'], [0.5, 2.0, 2.0], r_ids, r_dists)
-
-    # 5. 批量插入
-    index.batch_insert([
-        np.array([0.1, 0.1, 0.1, 0.1]),
-        np.array([0.2, 0.2, 0.2, 0.2])
-    ], ["id5", "id6"])
-    q5 = np.array([0.1, 0.1, 0.1, 0.1])
-    r_ids, r_dists = index.search(q5, 6)
-    print_test_case("批量插入后检索", ['id5', 'id6', 'id1', 'id3', 'id4'], [0.0, 0.04, 0.34, 0.84, 0.84], r_ids[:5], r_dists[:5], 2)
-
-    # ==== 持久化保存 ====
-    print(colored("\n--- 保存索引到磁盘 ---", "yellow"))
-    index.store(root_dir)
-    print(colored(f"数据已保存到目录: {root_dir}", "yellow"))
-
-    # ==== 内存对象清空 ====
-    del index
-    print(colored("内存对象已清除。", "yellow"))
-
-    # ==== 读取并检索 ====
-    user_input = input(colored("输入 yes 加载刚才保存的数据: ", "yellow"))
-    if user_input.strip().lower() == "yes":
-        index2 = FaissIndex.load(index_name, root_dir)
-        print(colored("数据已从磁盘恢复！", "green"))
-
-        r_ids, r_dists = index2.search(np.array([0.1, 0.1, 0.1, 0.1]), 5)
-        print_test_case("恢复后检索", ["id5", "id6", "id1", "id3", "id4"], [0.0, 0.04, 0.34, 0.84, 0.84], r_ids, r_dists, 2)
-    else:
-        print(colored("跳过加载测试。", "yellow"))
-
-    # ==== 清除磁盘数据 ====
-    user_input = input(colored("输入 yes 删除磁盘所有数据: ", "yellow"))
-    if user_input.strip().lower() == "yes":
-        shutil.rmtree(root_dir)
-        print(colored("所有数据已删除！", "green"))
-    else:
-        print(colored("未执行删除。", "yellow"))
-
-