isage-middleware 0.1.0__py3-none-any.whl

sage/middleware/services/vdb/vdb_service.py

@@ -0,0 +1,433 @@
+"""
+VDB Service - 向量数据库微服务
+提供向量存储和相似性搜索功能的服务任务，集成到SAGE DAG中
+使用嵌入式FAISS引擎，不依赖外部数据库
+"""
+from typing import Dict, Any, Optional, List, Union, Tuple, TYPE_CHECKING
+from dataclasses import dataclass
+import numpy as np
+import logging
+import time
+import uuid
+
+from sage.core.api.service.base_service import BaseService
+
+if TYPE_CHECKING:
+    from sage.core.factory.service_factory import ServiceFactory
+    from sage.kernel import ServiceContext
+
+
+@dataclass
+class VDBConfig:
+    """VDB服务配置"""
+    embedding_dimension: int = 384
+    index_type: str = "IndexFlatL2"  # FAISS索引类型
+    max_vectors: int = 1000000
+    similarity_threshold: float = 0.8
+    # FAISS索引配置
+    faiss_config: Dict[str, Any] = None
+
+
+@dataclass
+class VectorDocument:
+    """向量文档"""
+    id: str
+    vector: List[float]
+    metadata: Dict[str, Any]
+    text: Optional[str] = None
+    created_at: float = None
+    
+    def __post_init__(self):
+        if self.created_at is None:
+            self.created_at = time.time()
+
+
+class FaissVDBBackend:
+    """基于FAISS的向量数据库后端"""
+    
+    def __init__(self, embedding_dimension: int = 384, index_type: str = "IndexFlatL2", 
+                 faiss_config: Optional[Dict[str, Any]] = None):
+        self.embedding_dimension = embedding_dimension
+        self.index_type = index_type
+        self.faiss_config = faiss_config or {}
+        self.logger = logging.getLogger(__name__)
+        
+        # 导入FaissIndex
+        try:
+            from sage.middleware.services.vdb.search_engine.faiss_index import FaissIndex
+            
+            # 初始化FAISS索引
+            self.faiss_index = FaissIndex(
+                name="vdb_index",
+                dim=embedding_dimension,
+                config={
+                    "index_type": index_type,
+                    **self.faiss_config
+                }
+            )
+            
+            # 文档存储 (id -> VectorDocument)
+            self.documents: Dict[str, VectorDocument] = {}
+            
+            self.logger.info(f"FaissVDBBackend initialized with {index_type}, dim={embedding_dimension}")
+            
+        except ImportError as e:
+            self.logger.error(f"Failed to import FaissIndex: {e}")
+            raise ImportError("FAISS dependencies not available. Install with: pip install faiss-cpu")
+    
+    def add_vectors(self, vectors: List[VectorDocument]) -> List[str]:
+        """添加向量文档"""
+        try:
+            if not vectors:
+                return []
+            
+            # 准备向量和ID
+            vector_arrays = []
+            vector_ids = []
+            
+            for doc in vectors:
+                # 如果没有ID，生成一个
+                if not doc.id:
+                    doc.id = str(uuid.uuid4())
+                
+                # 验证向量维度
+                if len(doc.vector) != self.embedding_dimension:
+                    self.logger.warning(f"Vector dimension mismatch: expected {self.embedding_dimension}, got {len(doc.vector)}")
+                    continue
+                
+                vector_arrays.append(np.array(doc.vector, dtype=np.float32))
+                vector_ids.append(doc.id)
+                
+                # 存储文档
+                self.documents[doc.id] = doc
+            
+            if vector_arrays:
+                # 添加到FAISS索引
+                vectors_np = np.vstack(vector_arrays)
+                self.faiss_index.add(vectors_np, vector_ids)
+                
+                self.logger.debug(f"Added {len(vector_ids)} vectors to FAISS index")
+                return vector_ids
+            else:
+                return []
+                
+        except Exception as e:
+            self.logger.error(f"Error adding vectors: {e}")
+            return []
+    
+    def search_vectors(self, query_vector: List[float], top_k: int = 5,
+                      metadata_filter: Optional[Dict[str, Any]] = None) -> List[Tuple[VectorDocument, float]]:
+        """搜索相似向量"""
+        try:
+            if len(query_vector) != self.embedding_dimension:
+                raise ValueError(f"Query vector dimension mismatch: expected {self.embedding_dimension}, got {len(query_vector)}")
+            
+            # 转换查询向量
+            query_np = np.array([query_vector], dtype=np.float32)
+            
+            # 在FAISS中搜索
+            distances, indices, ids = self.faiss_index.search(query_np, top_k)
+            
+            results = []
+            for i, (distance, doc_id) in enumerate(zip(distances[0], ids[0])):
+                if doc_id in self.documents:
+                    doc = self.documents[doc_id]
+                    
+                    # 应用元数据过滤
+                    if metadata_filter and not self._matches_filter(doc.metadata, metadata_filter):
+                        continue
+                    
+                    results.append((doc, float(distance)))
+                else:
+                    self.logger.warning(f"Document {doc_id} found in index but not in storage")
+            
+            self.logger.debug(f"Search returned {len(results)} results")
+            return results
+            
+        except Exception as e:
+            self.logger.error(f"Error searching vectors: {e}")
+            return []
+    
+    def get_vector(self, doc_id: str) -> Optional[VectorDocument]:
+        """获取向量文档"""
+        return self.documents.get(doc_id)
+    
+    def delete_vectors(self, doc_ids: List[str]) -> int:
+        """删除向量文档"""
+        try:
+            deleted_count = 0
+            valid_ids = []
+            
+            for doc_id in doc_ids:
+                if doc_id in self.documents:
+                    del self.documents[doc_id]
+                    valid_ids.append(doc_id)
+                    deleted_count += 1
+            
+            # 从FAISS索引中删除
+            if valid_ids:
+                self.faiss_index.delete(valid_ids)
+            
+            self.logger.debug(f"Deleted {deleted_count} vectors")
+            return deleted_count
+            
+        except Exception as e:
+            self.logger.error(f"Error deleting vectors: {e}")
+            return 0
+    
+    def list_vectors(self, metadata_filter: Optional[Dict[str, Any]] = None) -> List[VectorDocument]:
+        """列出向量文档"""
+        if metadata_filter is None:
+            return list(self.documents.values())
+        
+        return [doc for doc in self.documents.values() 
+                if self._matches_filter(doc.metadata, metadata_filter)]
+    
+    def count(self) -> int:
+        """获取向量数量"""
+        return len(self.documents)
+    
+    def clear(self) -> None:
+        """清空所有向量"""
+        self.documents.clear()
+        # 重新初始化FAISS索引
+        try:
+            from sage.middleware.services.vdb.search_engine.faiss_index import FaissIndex
+            self.faiss_index = FaissIndex(
+                name="vdb_index",
+                dim=self.embedding_dimension,
+                config={
+                    "index_type": self.index_type,
+                    **self.faiss_config
+                }
+            )
+        except Exception as e:
+            self.logger.error(f"Error reinitializing FAISS index: {e}")
+    
+    def save_index(self, path: str) -> bool:
+        """保存索引到磁盘"""
+        try:
+            self.faiss_index.save(path)
+            return True
+        except Exception as e:
+            self.logger.error(f"Error saving index: {e}")
+            return False
+    
+    def load_index(self, path: str) -> bool:
+        """从磁盘加载索引"""
+        try:
+            from sage.middleware.services.vdb.search_engine.faiss_index import FaissIndex
+            self.faiss_index = FaissIndex(
+                name="vdb_index", 
+                dim=self.embedding_dimension,
+                load_path=path
+            )
+            return True
+        except Exception as e:
+            self.logger.error(f"Error loading index: {e}")
+            return False
+    
+    def _matches_filter(self, metadata: Dict[str, Any], filter_dict: Dict[str, Any]) -> bool:
+        """检查元数据是否匹配过滤条件"""
+        for key, value in filter_dict.items():
+            if key not in metadata:
+                return False
+            if metadata[key] != value:
+                return False
+        return True
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """获取索引统计信息"""
+        try:
+            return {
+                "total_vectors": self.count(),
+                "index_type": self.index_type,
+                "embedding_dimension": self.embedding_dimension,
+                "faiss_index_ntotal": getattr(self.faiss_index.index, 'ntotal', 0) if self.faiss_index.index else 0,
+                "max_vectors": getattr(self, 'max_vectors', -1)
+            }
+        except Exception as e:
+            self.logger.error(f"Error getting stats: {e}")
+            return {"error": str(e)}
+
+
+class VDBService(BaseService):
+    """
+    VDB服务任务
+    
+    提供向量存储和相似性搜索功能，可以在SAGE DAG中作为服务节点使用
+    使用嵌入式FAISS引擎，高性能的本地向量检索
+    """
+    
+    def __init__(self, service_factory: 'ServiceFactory', ctx: 'ServiceContext' = None):
+        super().__init__(service_factory, ctx)
+        
+        # 从service_factory获取配置
+        self.config: VDBConfig = getattr(service_factory, 'config', VDBConfig())
+        
+        # 初始化FAISS后端
+        self.backend = FaissVDBBackend(
+            embedding_dimension=self.config.embedding_dimension,
+            index_type=self.config.index_type,
+            faiss_config=self.config.faiss_config or {}
+        )
+        
+        self.logger.info(f"VDB Service '{self.service_name}' initialized with FAISS backend")
+        self.logger.info(f"Index type: {self.config.index_type}, Dimension: {self.config.embedding_dimension}")
+    
+    def _start_service_instance(self):
+        """启动VDB服务实例"""
+        self.logger.info(f"VDB Service '{self.service_name}' started")
+    
+    def _stop_service_instance(self):
+        """停止VDB服务实例"""
+        self.logger.info(f"VDB Service '{self.service_name}' stopped")
+    
+    # VDB操作方法 - 这些方法可以通过服务调用机制被调用
+    
+    def add_vectors(self, vectors: List[Dict[str, Any]]) -> List[str]:
+        """添加向量文档"""
+        self.logger.debug(f"Adding {len(vectors)} vectors")
+        
+        # 转换输入格式
+        docs = []
+        for v in vectors:
+            doc = VectorDocument(
+                id=v.get('id', ''),
+                vector=v['vector'],
+                metadata=v.get('metadata', {}),
+                text=v.get('text')
+            )
+            docs.append(doc)
+        
+        result = self.backend.add_vectors(docs)
+        self.logger.debug(f"Added {len(result)} vectors successfully")
+        return result
+    
+    def search_vectors(self, query_vector: List[float], top_k: int = 5,
+                      metadata_filter: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        """搜索相似向量"""
+        self.logger.debug(f"Searching vectors with top_k={top_k}")
+        
+        results = self.backend.search_vectors(query_vector, top_k, metadata_filter)
+        
+        # 转换输出格式
+        formatted_results = []
+        for doc, distance in results:
+            formatted_results.append({
+                'id': doc.id,
+                'vector': doc.vector,
+                'metadata': doc.metadata,
+                'text': doc.text,
+                'distance': distance,
+                'created_at': doc.created_at
+            })
+        
+        self.logger.debug(f"Search returned {len(formatted_results)} results")
+        return formatted_results
+    
+    def get_vector(self, doc_id: str) -> Optional[Dict[str, Any]]:
+        """获取向量文档"""
+        self.logger.debug(f"Getting vector: {doc_id}")
+        
+        doc = self.backend.get_vector(doc_id)
+        if doc:
+            return {
+                'id': doc.id,
+                'vector': doc.vector,
+                'metadata': doc.metadata,
+                'text': doc.text,
+                'created_at': doc.created_at
+            }
+        return None
+    
+    def delete_vectors(self, doc_ids: List[str]) -> int:
+        """删除向量文档"""
+        self.logger.debug(f"Deleting {len(doc_ids)} vectors")
+        result = self.backend.delete_vectors(doc_ids)
+        self.logger.debug(f"Deleted {result} vectors successfully")
+        return result
+    
+    def list_vectors(self, metadata_filter: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        """列出向量文档"""
+        self.logger.debug("Listing vectors")
+        
+        docs = self.backend.list_vectors(metadata_filter)
+        results = []
+        for doc in docs:
+            results.append({
+                'id': doc.id,
+                'vector': doc.vector,
+                'metadata': doc.metadata,
+                'text': doc.text,
+                'created_at': doc.created_at
+            })
+        
+        self.logger.debug(f"Listed {len(results)} vectors")
+        return results
+    
+    def count(self) -> int:
+        """获取向量数量"""
+        result = self.backend.count()
+        self.logger.debug(f"Vector count: {result}")
+        return result
+    
+    def clear(self) -> None:
+        """清空所有向量"""
+        self.logger.debug("Clearing all vectors")
+        self.backend.clear()
+    
+    def save_index(self, path: str) -> bool:
+        """保存索引到磁盘"""
+        self.logger.debug(f"Saving index to: {path}")
+        return self.backend.save_index(path)
+    
+    def load_index(self, path: str) -> bool:
+        """从磁盘加载索引"""
+        self.logger.debug(f"Loading index from: {path}")
+        return self.backend.load_index(path)
+    
+    def stats(self) -> Dict[str, Any]:
+        """获取服务统计信息"""
+        base_stats = self.get_statistics()
+        vdb_stats = self.backend.get_stats()
+        base_stats.update(vdb_stats)
+        return base_stats
+
+
+# 工厂函数，用于在DAG中创建VDB服务
+def create_vdb_service_factory(
+    service_name: str = "vdb_service",
+    embedding_dimension: int = 384,
+    index_type: str = "IndexFlatL2",
+    max_vectors: int = 1000000,
+    similarity_threshold: float = 0.8,
+    faiss_config: Optional[Dict[str, Any]] = None
+):
+    """
+    创建VDB服务工厂
+    
+    Args:
+        service_name: 服务名称
+        embedding_dimension: 向量维度
+        index_type: FAISS索引类型 ("IndexFlatL2", "IndexHNSWFlat", "IndexIVFFlat", etc.)
+        max_vectors: 最大向量数量
+        similarity_threshold: 相似度阈值
+        faiss_config: FAISS索引配置
+    
+    Returns:
+        ServiceFactory: 可以用于注册到环境的服务工厂
+    """
+    from sage.core.factory.service_factory import ServiceFactory
+    
+    config = VDBConfig(
+        embedding_dimension=embedding_dimension,
+        index_type=index_type,
+        max_vectors=max_vectors,
+        similarity_threshold=similarity_threshold,
+        faiss_config=faiss_config or {}
+    )
+    
+    factory = ServiceFactory(service_name, VDBService)
+    factory.config = config
+    return factory

sage/middleware/utils/__init__.py

@@ -0,0 +1,5 @@
+"""
+Llm module for SAGE framework
+"""
+
+__version__ = "1.0.0"

sage/middleware/utils/embedding/__init__.py

@@ -0,0 +1,35 @@
+"""
+Embedding methods and models for SAGE middleware.
+"""
+
+from .embedding_api import apply_embedding_model
+from .embedding_model import EmbeddingModel
+
+# Import specific embedding implementations
+from . import hf
+from . import ollama
+from . import siliconcloud
+from . import openai
+from . import bedrock
+from . import zhipu
+from . import mockembedder
+from . import _cohere
+from . import nvidia_openai
+from . import lollms
+from . import jina
+
+__all__ = [
+    'apply_embedding_model',
+    'EmbeddingModel',
+    'hf',
+    'ollama', 
+    'siliconcloud',
+    'openai',
+    'bedrock',
+    'zhipu',
+    'mockembedder',
+    '_cohere',
+    'nvidia_openai',
+    'lollms',
+    'jina'
+]

sage/middleware/utils/embedding/_cohere.py

@@ -0,0 +1,68 @@
+import os
+import cohere
+import asyncio
+
+import numpy as np
+
+
+async def cohere_embed(
+        texts: list[str], api_key: str, model: str = "embed-multilingual-v3.0", input_type: str = "classification",
+        embedding_types: list[str] = ["float"]
+) -> list[float]:
+    if api_key is None:
+        api_key = os.environ.get("COHERE_API_KEY")
+    # print(api_key)
+    co = cohere.AsyncClient(api_key=api_key)
+
+    response = await co.embed(
+        texts=texts,
+        model=model,
+        input_type=input_type,
+        # embedding_types=embedding_types
+    )
+    return response.embeddings
+
+import os
+import cohere
+
+def cohere_embed_sync(
+    texts: list[str],
+    api_key: str = None,
+    model: str = "embed-multilingual-v3.0",
+    input_type: str = "classification",
+    embedding_types: list[str] = ["float"]
+) -> list[list[float]]:
+    """
+    同步版本：使用 Cohere 同步客户端生成文本 embeddings。
+
+    Args:
+        texts: 文本列表
+        api_key: Cohere API Key
+        model: 模型名称
+        input_type: 输入类型，如 classification、search_document 等
+        embedding_types: 嵌入格式（默认 float）
+
+    Returns:
+        list[list[float]]: 每个文本对应的嵌入向量
+    """
+    if api_key is None:
+        api_key = os.environ.get("COHERE_API_KEY")
+    if api_key is None:
+        raise ValueError("Cohere API key must be provided.")
+
+    co = cohere.Client(api_key=api_key)
+
+    response = co.embed(
+        texts=texts,
+        model=model,
+        input_type=input_type,
+        embedding_types=embedding_types,
+    )
+    return response.embeddings
+
+
+
+
+
+
+