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

sage/middleware/components/neuromem/search_engine/kv_index/__init__.py

@@ -1,37 +1,135 @@
-"""
-SAGE - Streaming-Augmented Generative Execution
-"""
-
-# 动态版本加载
-def _load_version():
-    """从项目根目录动态加载版本信息"""
-    from pathlib import Path
-    
-    # 获取项目根目录
-    current_file = Path(__file__).resolve()
-    root_dir = current_file.parent.parent.parent.parent.parent.parent.parent.parent.parent
-    version_file = root_dir / "_version.py"
-    
-    # 加载版本信息
-    if version_file.exists():
-        version_globals = {}
-        with open(version_file, 'r', encoding='utf-8') as f:
-            exec(f.read(), version_globals)
-        return {
-            'version': version_globals.get('__version__', '0.1.4'),
-            'author': version_globals.get('__author__', 'SAGE Team'),
-            'email': version_globals.get('__email__', 'shuhao_zhang@hust.edu.cn')
-        }
-    
-    # 默认值
-    return {
-        'version': '0.1.4',
-        'author': 'SAGE Team', 
-        'email': 'shuhao_zhang@hust.edu.cn'
+# file sage/middleware/services/neuromem/search_engine/kv_index/__init__.py
+
+from typing import Any, Dict, List, Optional
+
+from .base_kv_index import BaseKVIndex
+from .bm25s_index import BM25sIndex
+
+
+class KVIndexFactory:
+    """
+    键值索引工厂类，用于创建和加载不同类型的索引
+    Key-Value Index Factory for creating and loading different types of indexes
+    """
+
+    # 支持的索引类型映射
+    _index_registry = {
+        "bm25s": BM25sIndex,
     }
 
-# 加载信息
-_info = _load_version()
-__version__ = _info['version']
-__author__ = _info['author']
-__email__ = _info['email']
+    @classmethod
+    def create_index(
+        cls,
+        index_type: str,
+        index_name: str,
+        texts: List[str],
+        ids: List[str],
+        config: Optional[Dict[str, Any]] = None,
+    ) -> BaseKVIndex:
+        """
+        创建指定类型的索引实例
+        Create an index instance of the specified type
+
+        Args:
+            index_type: 索引类型，如 "bm25s"
+            index_name: 索引名称
+            texts: 文本列表
+            ids: ID列表
+            config: 额外配置参数
+
+        Returns:
+            BaseKVIndex: 创建的索引实例
+
+        Raises:
+            ValueError: 当索引类型不支持时
+        """
+        if index_type not in cls._index_registry:
+            supported_types = list(cls._index_registry.keys())
+            raise ValueError(
+                f"Index type '{index_type}' not supported. Supported types: {supported_types}"
+            )
+
+        # 准备配置
+        index_config = config or {}
+        index_config["name"] = index_name
+
+        # 获取索引类并创建实例
+        index_class = cls._index_registry[index_type]
+        index = index_class(config=index_config)
+
+        # 构建索引
+        if texts and ids:
+            index.build_index(texts, ids)
+
+        return index
+
+    @classmethod
+    def load_index(cls, index_type: str, index_name: str, dir_path: str) -> BaseKVIndex:
+        """
+        从磁盘加载指定类型的索引实例
+        Load an index instance of the specified type from disk
+
+        Args:
+            index_type: 索引类型，如 "bm25s"
+            index_name: 索引名称
+            dir_path: 索引存储目录路径
+
+        Returns:
+            BaseKVIndex: 加载的索引实例
+
+        Raises:
+            ValueError: 当索引类型不支持时
+        """
+        if index_type not in cls._index_registry:
+            supported_types = list(cls._index_registry.keys())
+            raise ValueError(
+                f"Index type '{index_type}' not supported. Supported types: {supported_types}"
+            )
+
+        # 获取索引类并加载实例
+        index_class = cls._index_registry[index_type]
+        return index_class.load(index_name, dir_path)
+
+    @classmethod
+    def get_supported_types(cls) -> List[str]:
+        """
+        获取所有支持的索引类型
+        Get all supported index types
+
+        Returns:
+            List[str]: 支持的索引类型列表
+        """
+        return list(cls._index_registry.keys())
+
+    @classmethod
+    def register_index_type(cls, index_type: str, index_class: type):
+        """
+        注册新的索引类型
+        Register a new index type
+
+        Args:
+            index_type: 索引类型名称
+            index_class: 索引类
+        """
+        if not issubclass(index_class, BaseKVIndex):
+            raise ValueError("Index class must be a subclass of BaseKVIndex")
+
+        cls._index_registry[index_type] = index_class
+
+    @classmethod
+    def is_supported(cls, index_type: str) -> bool:
+        """
+        检查是否支持指定的索引类型
+        Check if the specified index type is supported
+
+        Args:
+            index_type: 索引类型
+
+        Returns:
+            bool: 是否支持
+        """
+        return index_type in cls._index_registry
+
+
+# 导出主要类和工厂
+__all__ = ["BaseKVIndex", "BM25sIndex", "KVIndexFactory"]

sage/middleware/components/neuromem/search_engine/kv_index/base_kv_index.py

@@ -1,14 +1,15 @@
 # file sage/core/sage.middleware.services.neuromem./search_engine/kv_index/base_kv_index.py
 
 from abc import ABC, abstractmethod
-from typing import List, Any, Dict, Optional
+from typing import Any, Dict, List, Optional
+
 
 class BaseKVIndex(ABC):
     def __init__(
         self,
         config: Optional[dict] = None,
         texts: Optional[List[str]] = None,
-        ids: Optional[List[str]] = None
+        ids: Optional[List[str]] = None,
     ):
         """
         初始化索引基类。

sage/middleware/components/neuromem/search_engine/kv_index/bm25s_index.py

@@ -2,12 +2,14 @@
 # python -m sage.core.sage.middleware.services.neuromem..search_engine.kv_index.bm25s_index
 
 import os
-import bm25s
 import shutil
-import Stemmer
-from typing import List, Optional, Dict, Any
+from typing import Any, Dict, List, Optional
 
-from sage.middleware.components.neuromem.search_engine.kv_index.base_kv_index import BaseKVIndex
+import bm25s
+import Stemmer
+from sage.middleware.components.neuromem.search_engine.kv_index.base_kv_index import (
+    BaseKVIndex,
+)
 
 
 class BM25sIndex(BaseKVIndex):
@@ -20,25 +22,25 @@ class BM25sIndex(BaseKVIndex):
         支持两种初始化方式：
         1. 直接通过声明来创建：传入 config，然后使用 build_index() 方法来构建索引
         2. 通过 BM25sIndex.load() 来加载：调用load方法
-        
+
         Initialize the BM25sIndex instance with two initialization methods:
         1. Direct creation: pass config, then use build_index() method to build the index
         2. Load from disk: use load method
         """
         self.config = config or {}
-        
+
         # 从config中获取必要参数，否则使用默认值
         self.name = self.config.get("name", None)
         if self.name is None:
             raise ValueError("索引名称(name)未在config中指定")
-        
+
         # 初始化基本属性
         self.ids: List[str] = []
         self.texts: List[str] = []
         self.tokens: List[List[str]] = []
         self.tokenizer = None
         self.bm25 = None
-        
+
         # BM25s 特定配置参数
         self.backend = self.config.get("backend", "numba")
         self.language = self.config.get("language", "auto")  # auto, zh, en
@@ -56,16 +58,13 @@ class BM25sIndex(BaseKVIndex):
             zh_flag = False
         else:  # auto
             zh_flag = self._is_chinese(texts[0]) if texts else False
-        
+
         if zh_flag:
-            return bm25s.tokenization.Tokenizer(
-                stopwords=self.custom_stopwords or 'zh'
-            )
+            return bm25s.tokenization.Tokenizer(stopwords=self.custom_stopwords or "zh")
         else:
             stemmer = Stemmer.Stemmer("english")
             return bm25s.tokenization.Tokenizer(
-                stopwords=self.custom_stopwords or 'en', 
-                stemmer=stemmer
+                stopwords=self.custom_stopwords or "en", stemmer=stemmer
             )
 
     def _build_index(self, texts: List[str], ids: List[str]):
@@ -76,7 +75,7 @@ class BM25sIndex(BaseKVIndex):
         self.ids = list(ids)
         self.texts = list(texts)
         self.tokenizer = self._get_tokenizer(self.texts)
-        self.tokens = self.tokenizer.tokenize(self.texts) # type: ignore
+        self.tokens = self.tokenizer.tokenize(self.texts)  # type: ignore
         self.bm25 = bm25s.BM25(corpus=self.texts, backend=self.backend)
         self.bm25.index(self.tokens)
 
@@ -97,7 +96,7 @@ class BM25sIndex(BaseKVIndex):
         if not self.texts:
             return
         self.tokenizer = self._get_tokenizer(self.texts)
-        self.tokens = self.tokenizer.tokenize(self.texts) # type: ignore
+        self.tokens = self.tokenizer.tokenize(self.texts)  # type: ignore
         self.bm25 = bm25s.BM25(corpus=self.texts, backend=self.backend)
         self.bm25.index(self.tokens)
 
@@ -106,7 +105,7 @@ class BM25sIndex(BaseKVIndex):
         判断字符串中是否包含中文字符。
         Detect whether the text contains Chinese characters.
         """
-        return any('\u4e00' <= ch <= '\u9fff' for ch in text)
+        return any("\u4e00" <= ch <= "\u9fff" for ch in text)
 
     def insert(self, text, doc_id):
         """
@@ -147,41 +146,41 @@ class BM25sIndex(BaseKVIndex):
         """
         if self.bm25 is None or len(self.ids) == 0:
             return []
-        query_token = self.tokenizer.tokenize([text])[0] # type: ignore
-        scores = self.bm25.get_scores(query_token) # type: ignore
+        query_token = self.tokenizer.tokenize([text])[0]  # type: ignore
+        scores = self.bm25.get_scores(query_token)  # type: ignore
         topk_idx = sorted(range(len(scores)), key=lambda i: -scores[i])[:topk]
         return [self.ids[i] for i in topk_idx]
 
-
     def store(self, dir_path: str) -> Dict[str, Any]:
         """
         将索引信息存储到指定目录，包含bm25模型、分词器、ids和texts。
         Store the index info into the specified directory, including bm25 model, tokenizer, ids, and texts.
         """
         os.makedirs(dir_path, exist_ok=True)
-        
+
         # 保存BM25模型
         if self.bm25 is not None:
-            self.bm25.vocab_dict = {str(k): v for k, v in self.bm25.vocab_dict.items()} # type: ignore
-            self.bm25.save(dir_path, corpus=None)# type: ignore
-        
+            self.bm25.vocab_dict = {str(k): v for k, v in self.bm25.vocab_dict.items()}  # type: ignore
+            self.bm25.save(dir_path, corpus=None)  # type: ignore
+
         # 保存分词器
         if self.tokenizer is not None:
-            self.tokenizer.save_vocab(dir_path) # type: ignore
-            self.tokenizer.save_stopwords(dir_path) # type: ignore
-        
+            self.tokenizer.save_vocab(dir_path)  # type: ignore
+            self.tokenizer.save_stopwords(dir_path)  # type: ignore
+
         # 保存配置信息
         meta = {
             "name": self.name,
             "backend": self.backend,
             "language": self.language,
             "custom_stopwords": self.custom_stopwords,
-            "config": self.config
+            "config": self.config,
         }
         with open(os.path.join(dir_path, "meta.json"), "w", encoding="utf-8") as f:
             import json
+
             json.dump(meta, f, ensure_ascii=False, indent=2)
-        
+
         # 保存ids和texts
         with open(os.path.join(dir_path, "ids.txt"), "w", encoding="utf-8") as f:
             for i in self.ids:
@@ -200,13 +199,14 @@ class BM25sIndex(BaseKVIndex):
         meta_path = os.path.join(dir_path, "meta.json")
         if os.path.exists(meta_path):
             import json
+
             with open(meta_path, "r", encoding="utf-8") as f:
                 meta = json.load(f)
             self.config = meta.get("config", {})
             self.backend = meta.get("backend", "numba")
             self.language = meta.get("language", "auto")
             self.custom_stopwords = meta.get("custom_stopwords", None)
-        
+
         # 加载BM25模型
         self.bm25 = bm25s.BM25.load(dir_path)
 
@@ -220,9 +220,9 @@ class BM25sIndex(BaseKVIndex):
             self.ids = [line.strip() for line in f.readlines()]
         with open(os.path.join(dir_path, "texts.txt"), "r", encoding="utf-8") as f:
             self.texts = [line.strip() for line in f.readlines()]
-            
+
         # 重建tokens
-        self.tokens = [self.tokenizer.tokenize([t], return_as="tuple")[0][0] for t in self.texts] # type: ignore
+        self.tokens = [self.tokenizer.tokenize([t], return_as="tuple")[0][0] for t in self.texts]  # type: ignore
         self.bm25.index(self.tokens)
 
     @classmethod
@@ -258,7 +258,7 @@ if __name__ == "__main__":
     texts = [
         "The quick brown fox jumps over the lazy dog.",
         "Hello world! This is a operator_test document.",
-        "Python is a great programming language."
+        "Python is a great programming language.",
     ]
     root_path = "./tmp_bm25_test"  # 用临时目录避免误删业务数据
     index_dir = root_path
@@ -270,7 +270,7 @@ if __name__ == "__main__":
         "name": index_name,
         "backend": "numba",
         "language": "auto",  # 可以指定 "zh", "en" 或 "auto"
-        "custom_stopwords": None
+        "custom_stopwords": None,
     }
     index = BM25sIndex(config=config)
     index.build_index(texts, ids)
@@ -317,4 +317,3 @@ if __name__ == "__main__":
     # 7. 清理测试目录
     print("\n== 清理测试目录 ==")
     BM25sIndex.clear(index_dir)
-

sage/middleware/components/neuromem/search_engine/vdb_index/__init__.py

@@ -1,37 +1,82 @@
-"""
-SAGE - Streaming-Augmented Generative Execution
-"""
-
-# 动态版本加载
-def _load_version():
-    """从项目根目录动态加载版本信息"""
-    from pathlib import Path
-    
-    # 获取项目根目录
-    current_file = Path(__file__).resolve()
-    root_dir = current_file.parent.parent.parent.parent.parent.parent.parent.parent.parent
-    version_file = root_dir / "_version.py"
-    
-    # 加载版本信息
-    if version_file.exists():
-        version_globals = {}
-        with open(version_file, 'r', encoding='utf-8') as f:
-            exec(f.read(), version_globals)
-        return {
-            'version': version_globals.get('__version__', '0.1.4'),
-            'author': version_globals.get('__author__', 'SAGE Team'),
-            'email': version_globals.get('__email__', 'shuhao_zhang@hust.edu.cn')
-        }
-    
-    # 默认值
-    return {
-        'version': '0.1.4',
-        'author': 'SAGE Team', 
-        'email': 'shuhao_zhang@hust.edu.cn'
-    }
-
-# 加载信息
-_info = _load_version()
-__version__ = _info['version']
-__author__ = _info['author']
-__email__ = _info['email']
+# file sage/middleware/services/neuromem/search_engine/vdb_index/__init__.py
+
+from typing import Any, Dict, Type  # noqa: F401
+
+from .base_vdb_index import BaseVDBIndex
+
+
+class VDBIndexFactory:
+    """向量数据库索引工厂类 - 简化版本"""
+
+    # 注册的索引类型映射
+    _index_registry: Dict[str, Type[BaseVDBIndex]] = {}
+
+    @classmethod
+    def register_index(cls, index_type: str, index_class: Type[BaseVDBIndex]):
+        """注册新的索引类型"""
+        if not issubclass(index_class, BaseVDBIndex):
+            raise TypeError(f"Index class {index_class} must inherit from BaseVDBIndex")
+        cls._index_registry[index_type.upper()] = index_class
+
+    def create_index(self, config: Dict[str, Any]) -> BaseVDBIndex:
+        """
+        创建索引 - 简化版本，只支持config方式创建
+
+        Args:
+            config: 配置字典，必须包含 name, dim, backend_type
+
+        Returns:
+            创建的索引实例
+        """
+        # 检查必要字段
+        required_fields = ["name", "dim"]
+        for field in required_fields:
+            if field not in config:
+                raise ValueError(f"配置中缺少必要字段: {field}")
+
+        backend_type = config.get("backend_type", "FAISS").upper()
+
+        if backend_type not in self._index_registry:
+            raise ValueError(f"不支持的索引类型: {backend_type}")
+
+        index_class = self._index_registry[backend_type]
+        return index_class(config=config)
+
+
+# 全局工厂实例
+index_factory = VDBIndexFactory()
+
+
+def register_index_type(index_type: str, index_class: Type[BaseVDBIndex]):
+    """注册新的索引类型"""
+    VDBIndexFactory.register_index(index_type, index_class)
+
+
+def create_index(config: Dict[str, Any]) -> BaseVDBIndex:
+    """创建索引"""
+    return index_factory.create_index(config)
+
+
+# 自动注册FAISS索引
+def _auto_register_indexes():
+    """自动注册已知的索引类型"""
+    try:
+        from .faiss_index import FaissIndex
+
+        register_index_type("FAISS", FaissIndex)
+    except ImportError:
+        pass
+
+
+# 执行自动注册
+_auto_register_indexes()
+
+
+# 导出公共接口
+__all__ = [
+    "VDBIndexFactory",
+    "BaseVDBIndex",
+    "index_factory",
+    "register_index_type",
+    "create_index",
+]

sage/middleware/components/neuromem/search_engine/vdb_index/base_vdb_index.py

@@ -1,13 +1,14 @@
 # file sage/core/sage.middleware.services.neuromem./search_engine/vdb_index/base_vdb_index.py
 
 from abc import ABC, abstractmethod
-from typing import List, Tuple, Dict, Any
+from typing import Any, Dict, List, Tuple
+
 import numpy as np
 
+
 class BaseVDBIndex(ABC):
     def __init__(self):
-        """
-        """
+        """ """
 
     @abstractmethod
     def insert(self, vector: np.ndarray, string_id: str) -> None:
@@ -30,10 +31,12 @@ class BaseVDBIndex(ABC):
         pass
 
     @abstractmethod
-    def search(self, query_vector: np.ndarray, topk: int = 10) -> Tuple[List[str], List[float]]:
+    def search(
+        self, query_vector: np.ndarray, topk: int = 10
+    ) -> Tuple[List[str], List[float]]:
         """向量检索，返回 (string_id, 距离) 列表"""
         pass
-    
+
     @classmethod
     @abstractmethod
     def load(cls, name: str, root_path: str) -> "BaseVDBIndex":
@@ -50,4 +53,3 @@ class BaseVDBIndex(ABC):
         Store the index data to the specified directory.
         """
         pass
-