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):
        """
        Initialize sync with data source configuration
        
        Args:
            config: Configuration for the data source
        """
        self.config = config
        self.file_parser = FileParser()
    
    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
            logger.error(f"Error getting file bytes for {file_path}: {e}")
            return None
    
    def _parse_file_content(self, file_bytes: bytes, file_path: str, host: Optional[str] = None) -> Optional[str]:
        """
        Parse file content using FileParser class
        
        Args:
            file_bytes: File content as bytes
            file_path: Path to the file (for getting file extension)
            host: Optional host address (for remote files, use the remote host address)
            
        Returns:
            Parsed file content if successfully parsed, None otherwise
        """
        try:
            # 使用 FileParser 解析文件内容
            # parse_file_content 方法的参数顺序是 content, filename
            documents = self.file_parser.parse_file_content(file_bytes, file_path, host=host)
            if documents:
                # 合并所有文档内容
                return '\n\n'.join(doc.text for doc in documents if doc.text)
            else:
                from loguru import logger
                logger.warning(f"No content extracted from {file_path}")
                return None
        except Exception as e:
            from loguru import logger
            logger.error(f"Error parsing file content for {file_path}: {e}")
            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 fetch_new_documents(self, last_sync_time=None, synced_doc_ids=None) -> List[Dict[str, Any]]:
        """
        Fetch new/updated documents from the data source since last sync time
        
        Args:
            last_sync_time: Last synchronization time
            synced_doc_ids: Set of document IDs that have already been synced
            
        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:
            # Get the appropriate sync class based on data source type
            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) -> type[BaseSync]:
    """
    Get the appropriate sync class based on data source type
    
    Args:
        source_type: Type of data source (database, local_folder, remote_folder)
        
    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.local_folder_sync import LocalFolderSync
    from sync.remote_folder_sync import RemoteFolderSync
    
    sync_classes = {
        'database': MySQLSync,  # Currently only MySQL, but can be extended
        'local_folder': LocalFolderSync,
        'remote_folder': RemoteFolderSync
    }
    
    if source_type not in sync_classes:
        raise ValueError(f"Unsupported data source type: {source_type}")
    
    return sync_classes[source_type]