RagFlow_Plugs/services/utils/infinity/__init__.py

"""
Infinity向量数据库主类
"""
from typing import List, Dict, Any, Optional
from services.utils.infinity.base import InfinityConnection
from services.utils.infinity.index import InfinityIndexManager
from services.utils.infinity.document import InfinityDocumentManager
from services.utils.infinity.search import InfinitySearchManager


class InfinityVectorDB:
    """
    Infinity向量数据库主类
    提供统一的接口，整合索引、文档和搜索功能
    """
    
    def __init__(self, host: str = None, port: int = None, user: str = None, password: str = None):
        """
        初始化Infinity向量数据库
        
        Args:
            host: Infinity主机地址
            port: Infinity端口
            user: 用户名
            password: 密码
        """
        # 初始化连接
        self.connection = InfinityConnection(host, port, user, password)
        
        # 初始化管理器
        self.index_manager = InfinityIndexManager(self.connection)
        self.document_manager = InfinityDocumentManager(self.connection)
        self.search_manager = InfinitySearchManager(self.connection)
    
    def create_index(self, index_name: str, mappings: Dict[str, Any] = None) -> bool:
        """创建索引"""
        return self.index_manager.create_index(index_name, mappings)
    
    def delete_index(self, index_name: str) -> bool:
        """删除索引"""
        return self.index_manager.delete_index(index_name)
    
    def index_exists(self, index_name: str) -> bool:
        """检查索引是否存在"""
        return self.index_manager.index_exists(index_name)
    
    def insert_document(self, index_name: str, document: Dict[str, Any], id: str = None) -> bool:
        """插入单个文档"""
        return self.document_manager.insert_document(index_name, document, id)
    
    def bulk_insert(self, index_name: str, documents: List[Dict[str, Any]]) -> Dict[str, Any]:
        """批量插入文档"""
        return self.document_manager.bulk_insert(index_name, documents)
    
    def update_document(self, index_name: str, document_id: str, document: Dict[str, Any]) -> bool:
        """更新单个文档"""
        return self.document_manager.update_document(index_name, document_id, document)
    
    def delete_document(self, index_name: str, document_id: str) -> bool:
        """删除单个文档"""
        return self.document_manager.delete_document(index_name, document_id)
    
    def get_document(self, index_name: str, document_id: str) -> Optional[Dict[str, Any]]:
        """获取单个文档"""
        return self.document_manager.get_document(index_name, document_id)
    
    def delete_by_query(self, index_name: str, query: Dict[str, Any]) -> Dict[str, Any]:
        """按查询条件删除文档"""
        return self.document_manager.delete_by_query(index_name, query)
    
    def search(self, index_name: str, query: Dict[str, Any], size: int = 10) -> Dict[str, Any]:
        """搜索文档"""
        return self.search_manager.search(index_name, query, size)
    
    def vector_search(self, index_name: str, vector_field: str, vector: List[float], size: int = 10, filter: Dict[str, Any] = None) -> Dict[str, Any]:
        """向量检索"""
        return self.search_manager.vector_search(index_name, vector_field, vector, size, filter)
    
    def hybrid_search(self, index_name: str, text_query: str, vector_field: str, vector: List[float], 
                     size: int = 10, text_weight: float = 0.5, vector_weight: float = 0.5) -> Dict[str, Any]:
        """混合检索"""
        return self.search_manager.hybrid_search(index_name, text_query, vector_field, vector, size, text_weight, vector_weight)
    
    def match_search(self, index_name: str, field: str, value: str, size: int = 10) -> Dict[str, Any]:
        """简单匹配搜索"""
        return self.search_manager.match_search(index_name, field, value, size)
    
    def match_all(self, index_name: str, size: int = 10) -> Dict[str, Any]:
        """匹配所有文档"""
        return self.search_manager.match_all(index_name, size)
    
    def close(self):
        """关闭连接"""
        self.connection.close()