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

sage/middleware/components/neuromem/micro_service/neuromem_vdb_service.py

@@ -0,0 +1,118 @@
+import os
+from typing import Dict, List, Any, Optional, Union
+from sage.core.api.service.base_service import BaseService
+from sage.common.utils.logging.custom_logger import CustomLogger
+from sage.middleware.components.neuromem.memory_manager import MemoryManager
+from sage.middleware.components.neuromem.memory_collection.vdb_collection import VDBMemoryCollection
+from sage.middleware.utils.embedding.embedding_api import apply_embedding_model
+
+
+class NeuroMemVDBService(BaseService):
+    def __init__(self, collection_name: Union[str, List[str]]):
+        self.manager = MemoryManager(self._get_default_data_dir())
+        self.online_register_collections: Dict[str, VDBMemoryCollection] = {}
+        
+        # 处理collection_name参数，支持单个字符串或字符串列表
+        if isinstance(collection_name, str):
+            collection_names = [collection_name]
+        else:
+            collection_names = collection_name
+            
+        # 连接已有的collection，不存在就会报错
+        for name in collection_names:
+            try:
+                collection = self.manager.get_collection(name)
+                if collection is None:
+                    raise ValueError(f"Collection '{name}' not found")
+                if not isinstance(collection, VDBMemoryCollection):
+                    raise TypeError(f"Collection '{name}' is not a VDBMemoryCollection")
+                
+                self.online_register_collections[name] = collection
+                self.logger.info(f"Successfully connected to collection: {name}")
+                
+                # 检查是否有global_index，没有就创建一个
+                if "global_index" not in collection.indexes:
+                    self.logger.info(f"Creating global_index for collection: {name}")
+                    collection.create_index("global_index", description="Global index for all data")
+                    
+            except Exception as e:
+                self.logger.error(f"Failed to connect to collection '{name}': {str(e)}")
+                raise
+
+    def retrieve(self, query_text: str, topk: int = 5, collection_name: Optional[str] = None, 
+                with_metadata: bool = False, **kwargs) -> List[Any]:
+        """
+        在所有online_register_collections上按照vdb_collection方式检索，默认使用global_index
+        
+        Args:
+            query_text: 查询文本
+            topk: 返回结果数量
+            collection_name: 指定collection名称，如果为None则在所有collection上检索
+            with_metadata: 是否返回metadata
+            **kwargs: 其他检索参数
+            
+        Returns:
+            检索结果列表
+        """
+        if not self.online_register_collections:
+            self.logger.warning("No collections are registered")
+            return []
+            
+        all_results = []
+        
+        # 如果指定了collection_name，只在该collection上检索
+        if collection_name:
+            if collection_name not in self.online_register_collections:
+                raise ValueError(f"Collection '{collection_name}' is not registered")
+            collections_to_search = {collection_name: self.online_register_collections[collection_name]}
+        else:
+            # 在所有注册的collection上检索
+            collections_to_search = self.online_register_collections
+            
+        for name, collection in collections_to_search.items():
+            try:
+                results = collection.retrieve(
+                    query_text, 
+                    topk=topk, 
+                    index_name="global_index",
+                    with_metadata=with_metadata,
+                    **kwargs
+                )
+                
+                # 为结果添加来源collection信息
+                if with_metadata:
+                    for result in results:
+                        if isinstance(result, dict):
+                            result['source_collection'] = name
+                        else:
+                            # 如果结果不是dict，转换为dict格式
+                            result = {'text': result, 'source_collection': name}
+                else:
+                    # 如果不要metadata，也可以选择添加来源信息
+                    results = [{'text': result, 'source_collection': name} for result in results]
+                    
+                all_results.extend(results)
+                self.logger.debug(f"Retrieved {len(results)} results from collection: {name}")
+                
+            except Exception as e:
+                self.logger.error(f"Error retrieving from collection '{name}': {str(e)}")
+                
+        # 如果有多个collection的结果，可以按相似度重新排序（这里简化处理）
+        return all_results[:topk] if len(all_results) > topk else all_results
+
+    def _create_index(self, collection_name: str, index_name: str, **kwargs):
+        """为指定collection创建索引"""
+        if collection_name not in self.online_register_collections:
+            raise ValueError(f"Collection '{collection_name}' is not registered")
+            
+        collection = self.online_register_collections[collection_name]
+        collection.create_index(index_name, **kwargs)
+        self.logger.info(f"Created index '{index_name}' for collection '{collection_name}'")
+
+    @classmethod
+    def _get_default_data_dir(cls):
+        """获取默认数据目录"""
+        cur_dir = os.getcwd()
+        data_dir = os.path.join(cur_dir, "data", "neuromem_vdb")
+        os.makedirs(data_dir, exist_ok=True)
+        return data_dir

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

@@ -0,0 +1,37 @@
+"""
+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
+    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']

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

@@ -0,0 +1,37 @@
+"""
+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']

sage/middleware/components/neuromem/search_engine/graph_index/base_graph_index.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from typing import Any, List
+
+class BaseGraphIndex(ABC):
+    def __init__(self, name: str):
+        """
+        基础图索引类初始化
+        :param name: 索引名称
+        """
+        self.name = name
+
+    @abstractmethod
+    def add_node(self, node_id: str, data: Any) -> None:
+        """添加一个节点"""
+        pass
+
+    @abstractmethod
+    def add_edge(self, from_node: str, to_node: str, weight: float = 1.0) -> None:
+        """添加一条有向边"""
+        pass
+
+    @abstractmethod
+    def remove_node(self, node_id: str) -> None:
+        """删除一个节点及其相关边"""
+        pass
+
+    @abstractmethod
+    def remove_edge(self, from_node: str, to_node: str) -> None:
+        """删除一条边"""
+        pass
+
+    @abstractmethod
+    def get_neighbors(self, node_id: str, k: int = 10) -> List[str]:
+        """获取指定节点的前 k 个邻居节点"""
+        pass
+
+    @abstractmethod
+    def has_node(self, node_id: str) -> bool:
+        """判断节点是否存在"""
+        pass

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

@@ -0,0 +1,37 @@
+"""
+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']

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

@@ -0,0 +1,37 @@
+"""
+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']

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

@@ -0,0 +1,76 @@
+# 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
+
+class BaseKVIndex(ABC):
+    def __init__(
+        self,
+        config: Optional[dict] = None,
+        texts: Optional[List[str]] = None,
+        ids: Optional[List[str]] = None
+    ):
+        """
+        初始化索引基类。
+        Initialize the base class for KV Index.
+        """
+        self.config = config or {}
+        self.name = self.config.get("name", None)
+
+    @abstractmethod
+    def insert(self, text: str, id: str) -> None:
+        """
+        插入一条新数据。
+        Insert a new entry.
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, id: str) -> None:
+        """
+        删除指定id的数据。
+        Delete an entry by id.
+        """
+        pass
+
+    @abstractmethod
+    def search(self, query: str, topk: int = 10) -> List[str]:
+        """
+        检索相关数据，返回最相关的id列表。
+        Search for relevant entries and return the most relevant ids.
+        """
+        pass
+
+    @abstractmethod
+    def update(self, id: str, new_text: str) -> None:
+        """
+        更新指定id的数据内容。
+        Update the entry corresponding to the given id.
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def load(cls, name: str, root_path: str) -> "BaseKVIndex":
+        """
+        加载索引实例。
+        Load the index instance.
+        """
+        pass
+
+    @abstractmethod
+    def store(self, root_path: str) -> Dict[str, Any]:
+        """
+        存储索引数据到指定目录。
+        Store the index data to the specified directory.
+        """
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def clear(dir_path: str) -> None:
+        """
+        删除指定目录下的所有索引数据。
+        Remove all index data under the specified directory.
+        """
+        pass