RAG/sync/base_sync.py

"""Base synchronization interface for all data sources"""
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Set, Optional
from config import BaseDataSourceConfig
from rag.file_parser import FileParser


class BaseSync(ABC):
    """Abstract base class for all data source syncers"""
    
    def __init__(self, config: BaseDataSourceConfig, vector_store_manager=None):
        """
        Initialize sync with data source configuration
        
        Args:
            config: Configuration for the data source
            vector_store_manager: Vector store manager for document existence checks
        """
        self.config = config
        self.file_parser = FileParser()
        self.vector_store_manager = vector_store_manager
    
    def _get_file_bytes(self, file_path: str, source_type: str, source_config: Optional[Dict[str, Any]] = None) -> Optional[bytes]:
        """
        Get file content as bytes based on source type
        
        Args:
            file_path: Path to the file
            source_type: Type of file source (api, filesystem, scp)
            source_config: Additional configuration for file source
            
        Returns:
            File content as bytes if successfully loaded, None otherwise
        """
        try:
            import requests
            
            if source_type == 'api':
                # 从 API 获取文件内容
                from config import settings
                url = f"{settings.FILE_DOWNLOAD_BASE_URL}?identifier={file_path}"
                response = requests.get(url, timeout=30)
                response.raise_for_status()
                return response.content
            elif source_type == 'filesystem':
                # 从本地文件系统获取文件内容
                base_path = source_config.get('file_system_base_path', '') if source_config else ''
                full_path = f"{base_path}/{file_path}" if base_path else file_path
                with open(full_path, 'rb') as f:
                    return f.read()
            elif source_type == 'scp':
                # 通过 SCP 获取文件内容
                import paramiko
                
                ssh_config = source_config or {}
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                
                # Connect to remote server
                if ssh_config.get('scp_key_path'):
                    # Use SSH key for authentication
                    private_key = paramiko.RSAKey.from_private_key_file(ssh_config['scp_key_path'])
                    ssh.connect(
                        hostname=ssh_config.get('scp_host', ''),
                        port=ssh_config.get('scp_port', 22),
                        username=ssh_config.get('scp_username', ''),
                        pkey=private_key
                    )
                else:
                    # Use password for authentication
                    ssh.connect(
                        hostname=ssh_config.get('scp_host', ''),
                        port=ssh_config.get('scp_port', 22),
                        username=ssh_config.get('scp_username', ''),
                        password=ssh_config.get('scp_password', '')
                    )
                
                # Use SFTP to download file as bytes
                with paramiko.SFTPClient.from_transport(ssh.get_transport()) as sftp:
                    with sftp.open(file_path, 'rb') as f:
                        content = f.read()
                
                ssh.close()
                return content
            else:
                from loguru import logger
                logger.error(f"Unsupported file source type: {source_type}")
                return None
        except Exception as e:
            from loguru import logger
            import traceback
            logger.error(f"Error getting file bytes for {file_path}: {e}")
            logger.error(f"Traceback:\n{traceback.format_exc()}")
            return None
    

    
    @abstractmethod
    def fetch_all_documents(self) -> List[Dict[str, Any]]:
        """
        Fetch all documents from the data source
        
        Returns:
            List of documents
        """
        pass
    
    @abstractmethod
    def doc_to_llamaindex_doc(self, doc: Dict) -> 'Document':
        """
        Convert data source document to LlamaIndex Document
        
        Args:
            doc: Document from the data source
            
        Returns:
            LlamaIndex Document object
        """
        pass
    
    def process_documents(self, docs: List[Dict]) -> List['Document']:
        """
        Process multiple documents into LlamaIndex Documents
        
        Args:
            docs: List of documents from the data source
            
        Returns:
            List of LlamaIndex Document objects
        """
        from loguru import logger
        import traceback
        documents = []
        for doc in docs:
            try:
                llamaindex_doc = self.doc_to_llamaindex_doc(doc)
                if len(llamaindex_doc.text.strip()) >= 10:
                    documents.append(llamaindex_doc)
                else:
                    logger.warning(f"跳过过短文档 (id: {llamaindex_doc.id_})，内容长度: {len(llamaindex_doc.text)} 字符")
            except Exception as e:
                doc_id = doc.get('id', 'unknown')
                logger.error(f"Error processing document {doc_id}: {e}")
                logger.error(f"Traceback:\n{traceback.format_exc()}")
                continue
        
        logger.info(f"Processed {len(documents)} documents")
        return documents
    
    @staticmethod
    def chunk_documents(documents: List['Document']) -> List['Document']:
        """
        Chunk documents into smaller pieces
        
        Args:
            documents: List of Document objects to chunk
            
        Returns:
            List of chunked Document objects
        """
        from llama_index.core import Document
        from llama_index.core.node_parser import SentenceSplitter
        from loguru import logger
        from config import settings
        
        node_parser = SentenceSplitter(
            chunk_size=settings.CHUNK_SIZE, #NOTE: 从settings中获取，默认1024
            chunk_overlap=settings.CHUNK_OVERLAP
        )
        
        chunked_docs = []
        for doc in documents:
            try:
                nodes = node_parser.get_nodes_from_documents([doc])
                # Get unique_doc_id from doc.id_ or metadata['doc_id']
                unique_doc_id = getattr(doc, 'id_', None) or (doc.metadata.get('doc_id') if doc.metadata else None)
                if not unique_doc_id:
                    # Fallback: use node.node_id if unique_doc_id is not available
                    unique_doc_id = f"doc_{id(doc)}"
                    logger.warning(f"Document has no id_ or doc_id in metadata, using fallback: {unique_doc_id}")
                
                # Convert nodes back to documents for storage
                for index, node in enumerate(nodes):
                    chunk_id = f"{unique_doc_id}_chunk_{index}"
                    chunked_doc = Document(
                        text=node.text,
                        id_=chunk_id,
                        metadata={
                            **doc.metadata,
                            "chunk_id": chunk_id,
                            "chunk_index": index,
                            "total_chunks": len(nodes),
                            "original_node_id": node.node_id
                        }
                    )
                    chunked_docs.append(chunked_doc)
            except Exception as e:
                # Access id_ property (not doc_id) for error logging
                doc_id = getattr(doc, 'id_', getattr(doc, 'doc_id', 'unknown'))
                import traceback
                logger.error(f"Error chunking document {doc_id}: {e}")
                logger.error(f"Traceback:\n{traceback.format_exc()}")
                continue
        
        logger.info(f"Chunked {len(documents)} documents into {len(chunked_docs)} chunks")
        return chunked_docs
    
    @abstractmethod
    def generate_doc_id(self, identifier: str) -> str:
        """
        Generate a unique document ID for different data sources
        
        Args:
            identifier: Unique identifier for the document (file path, record ID, etc.)
            
        Returns:
            Unique document ID
        """
        pass
    
    @abstractmethod
    def fetch_new_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
        """
        Fetch new/updated documents from the data source since last sync time
        
        Args:
            last_sync_time: Last synchronization time
            
        Returns:
            List of new/updated documents
        """
        pass
    
    @abstractmethod
    def get_synced_document_ids(self) -> Set[str]:
        """
        Get IDs of all synced documents from the data source
        
        Returns:
            Set of document IDs
        """
        pass
    
    @staticmethod
    @abstractmethod
    def check_data_source_exists(config: BaseDataSourceConfig) -> bool:
        """
        Check if the data source exists and is accessible
        
        Args:
            config: Data source configuration
            
        Returns:
            True if data source exists and is accessible, False otherwise
        """
        pass
    
    @staticmethod
    def check_data_sources_exist(configs: List[BaseDataSourceConfig]) -> tuple[bool, List[str]]:
        """
        Check if all data sources exist and are accessible
        
        Args:
            configs: List of data source configurations
            
        Returns:
            Tuple of (all_exist, missing_sources)
        """
        missing = []
        
        for config in configs:
            # 如果是数据库类型，使用db_type参数
            if config.type == 'database':
                sync_class = get_sync_class(config.type, getattr(config, 'db_type', 'mysql'))
            else:
                sync_class = get_sync_class(config.type)
            if not sync_class.check_data_source_exists(config):
                missing.append(config.name)
        
        return len(missing) == 0, missing


def get_sync_class(source_type: str, db_type: str = 'mysql') -> type[BaseSync]:
    """
    Get the appropriate sync class based on data source type

    Args:
        source_type: Type of data source (database, folder, git)
        db_type: Type of database (mysql, dameng, etc.) - only used when source_type is 'database'

    Returns:
        Sync class corresponding to the data source type

    Raises:
        ValueError: If source type is not supported
    """
    from sync.mysql_sync import MySQLSync
    from sync.folder_sync import FolderSync
    from sync.dameng_sync import DaMengSync
    from sync.git_sync import GitSync

    if source_type == 'database':
        # 根据数据库类型选择相应的同步类
        db_type_lower = db_type.lower()
        if db_type_lower == 'dameng':
            return DaMengSync
        else:  # 默认为mysql
            return MySQLSync
    elif source_type == 'folder':
        return FolderSync
    elif source_type == 'git':
        return GitSync
    else:
        raise ValueError(f"Unsupported data source type: {source_type}")