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

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

@@ -1,24 +1,31 @@
-from logging import Manager
 import os
-from re import M
-from typing import Dict, List, Any, Optional, Union
+from typing import Any, Dict, Optional
+
+from sage.middleware.components.neuromem.memory_collection.vdb_collection import (
+    VDBMemoryCollection,
+)
 from sage.middleware.components.neuromem.memory_manager import MemoryManager
-from sage.middleware.components.neuromem.memory_collection.vdb_collection import VDBMemoryCollection
+
 
 class NeuroMemVDB:
     def __init__(self):
         self.manager = MemoryManager(self._get_default_data_dir())
         self.online_register_collection: Dict[str, VDBMemoryCollection] = {}
 
-
-    def register_collection(self, collection_name: str, config = None):
+    def register_collection(self, collection_name: str, config=None):
         # 检查collection是否已存在
         existing_collections = self.manager.list_collection()
-        collection_exists = any(c['name'] == collection_name for c in existing_collections) if isinstance(existing_collections, list) else False
-        
+        collection_exists = (
+            any(c["name"] == collection_name for c in existing_collections)
+            if isinstance(existing_collections, list)
+            else False
+        )
+
         if collection_exists:
             if config is not None:
-                print(f"警告: Collection '{collection_name}' 已存在，忽略传入的config参数")
+                print(
+                    f"警告: Collection '{collection_name}' 已存在，忽略传入的config参数"
+                )
             # 正常连接并注册到online_register_collection
             collection = self.manager.get_collection(collection_name)
             if collection:
@@ -27,65 +34,87 @@ class NeuroMemVDB:
         else:
             if config is not None:
                 # 正常通过Manager创建并注册到online_register_collection
-                embedding_model = config.get('embedding_model')
-                dim = config.get('dim')
-                description = config.get('description', f"VDB collection: {collection_name}")
+                embedding_model = config.get("embedding_model")
+                dim = config.get("dim")
+                description = config.get(
+                    "description", f"VDB collection: {collection_name}"
+                )
                 collection = self.manager.create_collection(
                     name=collection_name,
                     backend_type="VDB",
                     description=description,
                     embedding_model=embedding_model,
-                    dim=dim
+                    dim=dim,
                 )
                 if collection:
                     self.online_register_collection[collection_name] = collection
                     print(f"成功创建新collection: {collection_name}")
             else:
                 # 警告 用默认参数创建一个collection并注册到online_register_collection
-                print(f"警告: Collection '{collection_name}' 不存在且未提供config，使用默认参数创建")
-                collection = self.manager.create_collection(
-                    name=collection_name,
-                    backend_type="VDB",
-                    description=f"Default VDB collection: {collection_name}"
+                print(
+                    f"警告: Collection '{collection_name}' 不存在且未提供config，使用默认参数创建"
                 )
+                default_config = {
+                    "name": collection_name,
+                    "backend_type": "VDB",
+                    "description": f"Default VDB collection: {collection_name}",
+                }
+                collection = self.manager.create_collection(default_config)
                 if collection:
                     self.online_register_collection[collection_name] = collection
                     print(f"成功创建默认collection: {collection_name}")
 
     def insert(self, raw_data: Any, metadata: Optional[Dict[str, Any]] = None):
-        # 该方法在所有 online_register_collection 均插入数据
+        # 该方法在所有 online_register_collection 均插入数据（仅存储，不创建索引）
         if not self.online_register_collection:
             print("警告: 没有注册的collection，无法插入数据")
             return
-        
+
         results = {}
         for collection_name, collection in self.online_register_collection.items():
             try:
-                stable_id = collection.insert(raw_data, metadata)
+                # 使用batch_insert_data方法插入单条数据（仅存储）
+                collection.batch_insert_data(
+                    [raw_data], [metadata] if metadata else None
+                )
+
+                # 生成stable_id用于返回
+                import hashlib
+                import json
+
+                key = raw_data
+                if metadata:
+                    key += json.dumps(metadata, sort_keys=True)
+                stable_id = hashlib.sha256(key.encode("utf-8")).hexdigest()
+
                 results[collection_name] = stable_id
-                print(f"成功在collection '{collection_name}' 中插入数据，ID: {stable_id}")
+                print(
+                    f"成功在collection '{collection_name}' 中插入数据，ID: {stable_id}"
+                )
             except Exception as e:
                 print(f"在collection '{collection_name}' 插入数据失败: {str(e)}")
                 results[collection_name] = None
         return results
-    
-    def retrieve(self, raw_data: Any, topk: int = 5):
+
+    def retrieve(self, raw_data: Any, topk: int = 5, index_name: str = "global_index"):
         # 测试类型的retrieve方法，不返回任何值，直接print出检索结果
         if not self.online_register_collection:
             print("警告: 没有注册的collection，无法检索数据")
             return
-        
+
         print(f"检索查询: {raw_data}")
         print("=" * 50)
-        
+
         for collection_name, collection in self.online_register_collection.items():
             try:
-                results = collection.retrieve(raw_data, topk=topk, with_metadata=True)
+                results = collection.retrieve(
+                    raw_data, index_name, topk=topk, with_metadata=True
+                )
                 print(f"Collection '{collection_name}' 检索结果:")
                 if results:
                     for i, result in enumerate(results, 1):
                         print(f"  {i}. 文本: {result['text']}")
-                        if result['metadata']:
+                        if result["metadata"]:
                             print(f"     元数据: {result['metadata']}")
                 else:
                     print("  无结果")
@@ -94,18 +123,35 @@ class NeuroMemVDB:
                 print(f"Collection '{collection_name}' 检索失败: {str(e)}")
                 print("-" * 30)
 
-    def build_index(self, index_name: str = "global_index", description: str = "全局索引"):
+    def build_index(
+        self, index_name: str = "global_index", description: str = "全局索引"
+    ):
         # 该方法在所有 online_register_collection 创建指定索引
         if not self.online_register_collection:
             print("警告: 没有注册的collection，无法创建索引")
             return
-        
+
         results = {}
         for collection_name, collection in self.online_register_collection.items():
             try:
-                collection.create_index(index_name, description=description)
-                results[collection_name] = True
-                print(f"成功在collection '{collection_name}' 中创建索引 '{index_name}'")
+                index_config = {
+                    "name": index_name,
+                    "description": description,
+                    "embedding_model": "default",  # 使用默认嵌入模型
+                    "dim": 384,  # 对应默认模型的维度
+                    "backend_type": "FAISS",  # 默认后端
+                }
+                success = collection.create_index(index_config)
+                if success:
+                    # 初始化索引，将现有数据插入索引
+                    collection.init_index(index_name)
+                    results[collection_name] = True
+                    print(
+                        f"成功在collection '{collection_name}' 中创建索引 '{index_name}'"
+                    )
+                else:
+                    results[collection_name] = False
+                    print(f"在collection '{collection_name}' 创建索引失败")
             except Exception as e:
                 print(f"在collection '{collection_name}' 创建索引失败: {str(e)}")
                 results[collection_name] = False
@@ -126,7 +172,7 @@ class NeuroMemVDB:
             if collection_name in self.online_register_collection:
                 del self.online_register_collection[collection_name]
                 print(f"已从在线注册列表中移除collection: {collection_name}")
-            
+
             # 通过manager删除collection（包括磁盘数据）
             self.manager.delete_collection(collection_name)
             print(f"成功清理collection: {collection_name}")
@@ -140,59 +186,60 @@ class NeuroMemVDB:
         data_dir = os.path.join(cur_dir, "data", "neuromem_vdb")
         os.makedirs(data_dir, exist_ok=True)
         return data_dir
-    
+
+
 if __name__ == "__main__":
     print("默认数据目录:", NeuroMemVDB._get_default_data_dir())
-    print("\n" + "="*60)
+    print("\n" + "=" * 60)
     print("开始测试 NeuroMemVDB")
-    print("="*60)
-    
+    print("=" * 60)
+
     # 创建NeuroMemVDB实例
     vdb = NeuroMemVDB()
-    
+
     # 测试1: 注册不存在的collection（无config，应该警告并使用默认参数）
     print("\n1. 测试注册不存在的collection（无config）")
     vdb.register_collection("test_collection1")
-    
+
     # 测试2: 注册不存在的collection（有config）
     print("\n2. 测试注册不存在的collection（有config）")
     config = {
-        'description': '测试用的VDB集合',
-        'embedding_model': None,  # 使用默认embedding模型
-        'dim': None  # 使用默认维度
+        "description": "测试用的VDB集合",
+        "embedding_model": None,  # 使用默认embedding模型
+        "dim": None,  # 使用默认维度
     }
     vdb.register_collection("test_collection2", config)
-    
+
     # 测试3: 插入数据
     print("\n3. 测试插入数据")
     vdb.insert("Python是一种编程语言", {"type": "test", "priority": "high"})
     vdb.insert("机器学习是人工智能的重要分支", {"type": "demo", "priority": "low"})
     vdb.insert("向量数据库用于存储和检索高维向量数据")
-    
+
     # 测试4: 创建索引
     print("\n4. 测试创建索引")
     vdb.build_index("custom_index", "自定义测试索引")
-    
+
     # 测试5: 检索数据
     print("\n5. 测试检索数据")
     vdb.retrieve("编程语言", topk=3)
-    
+
     # 测试6: 保存到磁盘
     print("\n6. 测试保存到磁盘")
     vdb.store_to_disk()
-    
+
     # 测试7: 重新注册已存在的collection（应该警告）
     print("\n7. 测试重新注册已存在的collection")
     vdb.register_collection("test_collection1", {"description": "这个config应该被忽略"})
-    
+
     # 测试8: 清理collection
     print("\n8. 测试清理collection")
     vdb.clean_collection("test_collection1")
-    
+
     # 测试9: 再次检索（应该只有test_collection2的结果）
     print("\n9. 清理后再次检索")
     vdb.retrieve("编程语言", topk=3)
-    
-    print("\n" + "="*60)
+
+    print("\n" + "=" * 60)
     print("测试完成")
-    print("="*60)
+    print("=" * 60)

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

@@ -1,23 +1,24 @@
 import os
-from typing import Dict, List, Any, Optional, Union
+from typing import Any, Dict, List, 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_collection.vdb_collection import (
+    VDBMemoryCollection,
+)
 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:
@@ -26,77 +27,94 @@ class NeuroMemVDBService(BaseService):
                     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")
-                    
+                    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]:
+    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]}
+            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, 
+                    query_text,
+                    topk=topk,
                     index_name="global_index",
                     with_metadata=with_metadata,
-                    **kwargs
+                    **kwargs,
                 )
-                
+
                 # 为结果添加来源collection信息
                 if with_metadata:
                     for result in results:
                         if isinstance(result, dict):
-                            result['source_collection'] = name
+                            result["source_collection"] = name
                         else:
                             # 如果结果不是dict，转换为dict格式
-                            result = {'text': result, 'source_collection': name}
+                            result = {"text": result, "source_collection": name}
                 else:
                     # 如果不要metadata，也可以选择添加来源信息
-                    results = [{'text': result, 'source_collection': name} for result in results]
-                    
+                    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}")
-                
+                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)}")
-                
+                self.logger.error(
+                    f"Error retrieving from collection '{name}': {str(e)}"
+                )
+
         # 如果有多个collection的结果，可以按相似度重新排序（这里简化处理）
         return all_results[:topk] if len(all_results) > topk else all_results
 
@@ -104,10 +122,12 @@ class NeuroMemVDBService(BaseService):
         """为指定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}'")
+        self.logger.info(
+            f"Created index '{index_name}' for collection '{collection_name}'"
+        )
 
     @classmethod
     def _get_default_data_dir(cls):

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

@@ -2,36 +2,11 @@
 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']
+# 直接从本包的_version模块加载版本信息
+try:
+    from sage.middleware._version import __author__, __email__, __version__
+except ImportError:
+    # 备用硬编码版本
+    __version__ = "0.1.4"
+    __author__ = "IntelliStream Team"
+    __email__ = "shuhao_zhang@hust.edu.cn"

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

@@ -2,36 +2,11 @@
 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']
+# 直接从本包的_version模块加载版本信息
+try:
+    from sage.middleware._version import __author__, __email__, __version__
+except ImportError:
+    # 备用硬编码版本
+    __version__ = "0.1.4"
+    __author__ = "IntelliStream Team"
+    __email__ = "shuhao_zhang@hust.edu.cn"

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

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, List
 
+
 class BaseGraphIndex(ABC):
     def __init__(self, name: str):
         """

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

@@ -0,0 +1,12 @@
+"""
+SAGE - Streaming-Augmented Generative Execution
+"""
+
+# 直接从本包的_version模块加载版本信息
+try:
+    from sage.middleware._version import __author__, __email__, __version__
+except ImportError:
+    # 备用硬编码版本
+    __version__ = "0.1.4"
+    __author__ = "IntelliStream Team"
+    __email__ = "shuhao_zhang@hust.edu.cn"