isage-middleware 0.1.0__py3-none-any.whl

sage/middleware/services/kv/search_engine/bm25s_index.py

@@ -0,0 +1,238 @@
+# file sage/core/sage.service.memory./search_engine/kv_index/bm25s_index.py
+# python -m sage.core.sage.service.memory..search_engine.kv_index.bm25s_index
+
+import os
+import bm25s
+import shutil
+import Stemmer
+from typing import List, Optional, Dict, Any
+
+from sage.middleware.services.kv.search_engine.base_kv_index import BaseKVIndex
+
+
+class BM25sIndex(BaseKVIndex):
+    def __init__(
+        self,
+        name: str,
+        texts: Optional[List[str]] = None,
+        ids: Optional[List[str]] = None,
+        load_path: Optional[str] = None,
+    ):
+        """
+        Initialize BM25sIndex.
+        支持两种初始化方式：传入文本和id新建索引，或从指定目录加载已有索引。
+        Supports two initialization modes: create a new index with texts and ids, or load an existing index from directory.
+        """
+        self.name = name
+        self.ids: List[str] = []
+        self.texts: List[str] = []
+        self.tokens: List[List[str]] = []
+        self.tokenizer = None
+        self.bm25 = None
+
+        if load_path is not None:
+            self._load(load_path)
+        elif texts is not None and ids is not None:
+            assert len(texts) == len(ids), "texts and ids must have the same length."
+            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.bm25 = bm25s.BM25(corpus=self.texts, backend="numba")
+            self.bm25.index(self.tokens)
+        else:
+            raise ValueError("Must provide either texts+ids or load_path.")
+
+    def _get_tokenizer(self, texts: List[str]):
+        """
+        根据文本内容选择合适的分词器（中文或英文）。
+        Select appropriate tokenizer (Chinese or English) according to the content of texts.
+        """
+        zh_flag = self._is_chinese(texts[0])
+        if zh_flag:
+            return bm25s.tokenization.Tokenizer(stopwords='zh')
+        else:
+            stemmer = Stemmer.Stemmer("english")
+            return bm25s.tokenization.Tokenizer(stopwords='en', stemmer=stemmer)
+
+    def _rebuild(self):
+        """
+        重新构建分词器、分词结果和BM25索引。
+        Rebuild the tokenizer, tokens, and BM25 index.
+        """
+        self.tokenizer = self._get_tokenizer(self.texts)
+        self.tokens = self.tokenizer.tokenize(self.texts) # type: ignore
+        self.bm25 = bm25s.BM25(corpus=self.texts, backend="numba")
+        self.bm25.index(self.tokens)
+
+    def _is_chinese(self, text: str):
+        """
+        判断字符串中是否包含中文字符。
+        Detect whether the text contains Chinese characters.
+        """
+        return any('\u4e00' <= ch <= '\u9fff' for ch in text)
+
+    def insert(self, text, doc_id):
+        """
+        插入新的文本和id，并重建索引。
+        Insert a new text and doc_id, then rebuild the index.
+        """
+        self.texts.append(text)
+        self.ids.append(doc_id)
+        self._rebuild()
+
+    def delete(self, id: str) -> None:
+        """
+        根据id删除对应的文本，并重建索引。
+        Delete the text corresponding to the given id, then rebuild the index.
+        """
+        if id not in self.ids:
+            return
+        idx = self.ids.index(id)
+        self.ids.pop(idx)
+        self.texts.pop(idx)
+        self._rebuild()
+
+    def update(self, id: str, new_text: str) -> None:
+        """
+        更新指定id的文本内容，并重建索引。
+        Update the text of the given id, then rebuild the index.
+        """
+        if id not in self.ids:
+            return
+        idx = self.ids.index(id)
+        self.texts[idx] = new_text
+        self._rebuild()
+
+    def search(self, text: str, topk: int = 5) -> List[str]:
+        """
+        对输入文本进行检索，返回最相关的topk个id。
+        Search for the most relevant texts and return the top-k ids.
+        """
+        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
+        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)
+        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.tokenizer.save_vocab(dir_path) # type: ignore
+        self.tokenizer.save_stopwords(dir_path) # type: ignore
+        with open(os.path.join(dir_path, "ids.txt"), "w", encoding="utf-8") as f:
+            for i in self.ids:
+                f.write(i + "\n")
+        with open(os.path.join(dir_path, "texts.txt"), "w", encoding="utf-8") as f:
+            for t in self.texts:
+                f.write(t.replace("\n", " ") + "\n")
+        return {"index_path": dir_path}
+
+    def _load(self, dir_path: str):
+        """
+        从目录加载索引及相关内容，包括bm25模型、分词器、ids和texts。
+        Load index and related data from directory, including bm25 model, tokenizer, ids, and texts.
+        """
+        self.bm25 = bm25s.BM25.load(dir_path)
+
+        self.tokenizer = bm25s.tokenization.Tokenizer()
+        self.tokenizer.load_vocab(dir_path)
+        self.tokenizer.load_stopwords(dir_path)
+
+        with open(os.path.join(dir_path, "ids.txt"), "r", encoding="utf-8") as f:
+            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()]
+        self.tokens = [self.tokenizer.tokenize([t], return_as="tuple")[0][0] for t in self.texts] # type: ignore
+        self.bm25.index(self.tokens)
+
+    @classmethod
+    def load(cls, name: str, dir_path: str) -> "BM25sIndex":
+        """
+        通过名称和根路径加载一个BM25sIndex实例。
+        Load a BM25sIndex instance by name and root path.
+        """
+        return cls(name=name, load_path=dir_path)
+
+    @staticmethod
+    def clear(dir_path: str):
+        """
+        删除指定名称下的所有索引数据。
+        Remove all index data under the specified name.
+        """
+        try:
+            shutil.rmtree(dir_path)
+            print(f"Cleared: {dir_path}")
+        except FileNotFoundError:
+            print(f"Directory does not exist, nothing to clear: {dir_path}")
+        except Exception as e:
+            print(f"Failed to clear: {e}")
+
+
+if __name__ == "__main__":
+    # 简单数据
+    ids = ["a", "b", "c"]
+    texts = [
+        "The quick brown fox jumps over the lazy dog.",
+        "Hello world! This is a operator_test document.",
+        "Python is a great programming language."
+    ]
+    root_path = "./tmp_bm25_test"  # 用临时目录避免误删业务数据
+    index_dir = root_path
+    index_name = "demo"
+
+    # 1. 初始化索引
+    print("\n== 初始化并首检 ==")
+    index = BM25sIndex(name=index_name, texts=texts, ids=ids)
+    print("初始检索 'Python':", index.search("Python"))
+    print("初始检索 'hello':", index.search("hello"))
+    print("初始检索 'fox':", index.search("fox"))
+
+    # 2. 插入新文档后检索
+    print("\n== 插入新文档 ==")
+    index.insert("This document mentions python and fox together.", "d")
+    print("插入后检索 'python':", index.search("python"))
+    print("插入后检索 'fox':", index.search("fox"))
+
+    # 3. 删除文档后检索
+    print("\n== 删除文档 ==")
+    index.delete("b")
+    print("删除 'b' 后检索 'hello':", index.search("hello"))
+    print("删除 'b' 后检索 'operator_test':", index.search("operator_test"))
+    print("删除 'b' 后检索 'python':", index.search("python"))
+
+    # 4. 更新文档后检索
+    print("\n== 更新文档 ==")
+    index.update("c", "Hello world! Now c document talks about foxes.")
+    print("更新 'c' 后检索 'fox':", index.search("fox"))
+    print("更新 'c' 后检索 'python':", index.search("python"))
+
+    # 5. 保存索引
+    print("\n== 保存索引到磁盘 ==")
+    store_info = index.store(index_dir)
+    print("索引保存路径:", store_info["index_path"])
+
+    # 6. 等待用户输入 'yes' 后加载索引并检索
+    print("\n== 测试持久化（请手动输入 yes 继续）==")
+    user_input = input("输入 'yes' 以继续测试 load 并检索：")
+    if user_input.strip().lower() == "yes":
+        index_loaded = BM25sIndex.load(name=index_name, dir_path=index_dir)
+        print("持久化load后检索 'fox':", index_loaded.search("fox"))
+        print("持久化load后检索 'python':", index_loaded.search("python"))
+        print("持久化load后检索 'hello':", index_loaded.search("hello"))
+        print("ids序列:", index_loaded.ids)
+    else:
+        print("用户未输入 'yes'，测试提前结束。")
+
+    # 7. 清理测试目录
+    print("\n== 清理测试目录 ==")
+    BM25sIndex.clear(index_dir)
+

sage/middleware/services/memory/__init__.py

@@ -0,0 +1,12 @@
+"""
+Memory Service Module
+记忆服务模块 - 提供统一的记忆管理和编排服务
+"""
+
+# 主要的记忆服务
+from .memory_service import MemoryService, create_memory_service_factory
+
+__all__ = [
+    "MemoryService", 
+    "create_memory_service_factory"
+]

sage/middleware/services/memory/examples/dag_microservices_demo.py

@@ -0,0 +1,220 @@
+"""
+SAGE 微服务架构使用示例
+展示如何在应用程序中注册和使用KV、VDB、Memory服务
+"""
+import asyncio
+import time
+from typing import List
+import numpy as np
+
+# 导入SAGE环境和服务
+from sage.core.api.local_environment import LocalEnvironment
+from sage.middleware.services import (
+    MemoryService,
+    create_kv_service_factory,
+    create_vdb_service_factory, 
+    create_memory_service_factory
+)
+
+
+class SampleApplication:
+    """示例应用程序，展示如何使用微服务架构"""
+    
+    def __init__(self):
+        # 创建SAGE环境
+        self.env = LocalEnvironment("microservices_demo", {})
+        
+    def setup_services(self):
+        """在应用中注册微服务"""
+        print("🔧 注册微服务到SAGE环境...")
+        
+        # 注册KV服务
+        kv_factory = create_kv_service_factory(
+            service_name="kv_service",
+            backend_type="memory",  # 使用内存后端
+            max_size=1000,
+            ttl_seconds=3600  # 1小时过期
+        )
+        self.env.register_service("kv_service", kv_factory.service_class, kv_factory)
+        
+        # 注册VDB服务
+        vdb_factory = create_vdb_service_factory(
+            service_name="vdb_service",
+            collection_name="demo_vectors",
+            dimension=384,
+            persist_directory="./demo_vectors"
+        )
+        self.env.register_service("vdb_service", vdb_factory.service_class, vdb_factory)
+        
+        # 注册Memory编排服务
+        memory_factory = create_memory_service_factory(
+            service_name="memory_service",
+            kv_service_name="kv_service",
+            vdb_service_name="vdb_service"
+        )
+        self.env.register_service("memory_service", memory_factory.service_class, memory_factory)
+        
+        print("✅ 所有服务已注册")
+    
+    def run_demo(self):
+        """运行演示"""
+        print("🚀 启动微服务演示")
+        print("=" * 50)
+        
+        # 设置服务
+        self.setup_services()
+        
+        # 创建一个简单的数据流来演示服务使用
+        data_stream = self.env.from_memory_source([
+            {"id": 1, "content": "用户询问了关于Python的问题", "session": "session_1"},
+            {"id": 2, "content": "AI助手回答了Python基础知识", "session": "session_1"},
+            {"id": 3, "content": "用户请求更多代码示例", "session": "session_1"},
+            {"id": 4, "content": "讨论了机器学习算法", "session": "session_2"},
+            {"id": 5, "content": "解释了神经网络原理", "session": "session_2"}
+        ])
+        
+        # 定义处理函数
+        def process_conversation(data):
+            """处理对话数据的函数"""
+            # 在这里我们可以使用服务调用
+            # 注意：在实际的SAGE函数中，可以通过 self.call_service 访问服务
+            
+            print(f"处理对话: {data['content'][:30]}...")
+            
+            # 模拟向量化（在实际应用中，这里会调用embedding服务）
+            content_vector = np.random.random(384).tolist()
+            
+            # 这里展示了服务调用的概念
+            # 在实际的SAGE函数中，代码会是这样：
+            # 
+            # # 存储到KV
+            # self.call_service["kv_service"].put(f"conv:{data['id']}", {
+            #     "content": data['content'],
+            #     "session": data['session'],
+            #     "timestamp": time.time()
+            # })
+            # 
+            # # 存储记忆
+            # memory_id = self.call_service["memory_service"].store_memory(
+            #     session_id=data['session'],
+            #     content=data['content'],
+            #     vector=content_vector,
+            #     memory_type="conversation"
+            # )
+            
+            return {
+                "processed": True,
+                "memory_id": f"mock_memory_{data['id']}",
+                "vector_dim": len(content_vector)
+            }
+        
+        # 应用处理函数
+        processed_stream = data_stream.map(process_conversation)
+        
+        # 执行并收集结果
+        print("\n📊 处理结果:")
+        results = processed_stream.collect()
+        
+        for i, result in enumerate(results, 1):
+            print(f"  {i}. ✅ 已处理 - Memory ID: {result['memory_id']}")
+        
+        print(f"\n🎯 总共处理了 {len(results)} 条对话记录")
+        
+        # 展示服务调用的概念
+        self.show_service_usage_concept()
+    
+    def show_service_usage_concept(self):
+        """展示服务使用概念"""
+        print("\n" + "=" * 50)
+        print("💡 在SAGE函数中使用服务的示例代码:")
+        print("=" * 50)
+        
+        example_code = '''
+# 在SAGE Function中使用微服务的示例
+
+class ConversationProcessor(BaseFunction):
+    """对话处理函数"""
+    
+    def process(self, conversation_data):
+        session_id = conversation_data['session_id']
+        content = conversation_data['content']
+        
+        # 1. 调用KV服务存储原始数据
+        kv_success = self.call_service["kv_service"].put(
+            f"raw:{session_id}", 
+            conversation_data
+        )
+        
+        # 2. 生成向量表示（假设有embedding服务）
+        vector = self.call_service["embedding_service"].encode(content)
+        
+        # 3. 调用Memory服务存储记忆
+        memory_id = self.call_service["memory_service"].store_memory(
+            session_id=session_id,
+            content=content,
+            vector=vector,
+            memory_type="conversation"
+        )
+        
+        # 4. 搜索相关历史记忆
+        related_memories = self.call_service["memory_service"].search_memories(
+            query_vector=vector,
+            session_id=session_id,
+            limit=5
+        )
+        
+        return {
+            "memory_id": memory_id,
+            "related_count": len(related_memories),
+            "kv_stored": kv_success
+        }
+
+# 在DAG中注册和使用
+def create_conversation_dag():
+    env = LocalEnvironment("conversation_app", {})
+    
+    # 注册微服务
+    env.register_service("kv_service", KVService, create_kv_service_factory())
+    env.register_service("vdb_service", VDBService, create_vdb_service_factory())
+    env.register_service("memory_service", MemoryService, create_memory_service_factory())
+    
+    # 创建数据流
+    stream = env.from_kafka_source(...)
+    
+    # 应用处理函数（自动访问服务）
+    processed = stream.map(ConversationProcessor())
+    
+    return processed
+        '''
+        
+        print(example_code)
+        print("\n" + "=" * 50)
+        print("🔍 关键概念:")
+        print("1. 服务作为Service Tasks在DAG中运行")
+        print("2. 函数通过 self.call_service[service_name] 调用服务")
+        print("3. 服务可以是本地任务或Ray分布式任务")
+        print("4. 应用程序控制服务的生命周期")
+        print("5. 服务间通过SAGE的队列机制通信")
+
+
+def main():
+    """主函数"""
+    app = SampleApplication()
+    
+    try:
+        app.run_demo()
+        print("\n✅ 演示完成!")
+        print("\n📖 查看更多信息:")
+        print("  - 微服务代码: packages/sage-middleware/src/sage/service/")
+        print("  - 使用指南: packages/sage-middleware/MICROSERVICES_GUIDE.md")
+        
+    except KeyboardInterrupt:
+        print("\n\n👋 演示被中断")
+    except Exception as e:
+        print(f"\n❌ 演示出错: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    main()