RAG/sync/folder_sync.py

"""Folder synchronization implementation for local and remote folders"""
import os
import re
from typing import List, Dict, Any, Set
from datetime import datetime
from pathlib import Path
from loguru import logger
from config import BaseDataSourceConfig
from sync.base_sync import BaseSync
from rag.file_parser import FileParser
from llama_index.core import Document
import paramiko


class SSHClient:
    """SSH client for testing connections to remote servers"""
    
    def __init__(self, host, port=22, username=None, password=None):
        """
        Initialize SSH client
        
        Args:
            host: Host address
            port: Port number
            username: Username
            password: Password
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self._client = None
    
    def test_connection(self):
        """
        Test SSH connection
        
        Returns:
            bool: True if connection is successful, False otherwise
        """
        try:
            # Create SSH client
            self._client = paramiko.SSHClient()
            self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # Connect with timeout
            self._client.connect(
                hostname=self.host,
                port=self.port,
                username=self.username,
                password=self.password,
                timeout=10,
                allow_agent=True,
                look_for_keys=False
            )
            
            # Connection successful
            return True
        finally:
            # Always close the connection
            if self._client:
                self._client.close()


class FolderSync(BaseSync):
    """Handle synchronization between folder (local or remote) and ChromaDB"""
    
    def __init__(self, config: BaseDataSourceConfig, vector_store_manager=None):
        """
        Initialize folder sync with configuration
        
        Args:
            config: Folder configuration
        """
        super().__init__(config)
        self.file_parser = FileParser()
        self._ssh_client = None
        self._sftp_client = None
        self.vector_store_manager = vector_store_manager
    
    def fetch_all_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
        """
        Fetch documents from the folder
        
        Args:
            last_sync_time: Last synchronization time (for incremental sync)
            
        Returns:
            List of documents
        """
        documents = []
        self._connect()
        try:
            files = self._get_all_files()
            
            for file_path in files:
                                
                # Parse file content
                try:
                    # 检查文件扩展名是否在支持的列表中
                    if Path(file_path).suffix.lower() not in FileParser.SUPPORTED_EXTENSIONS:
                        logger.debug(f"Skipping unsupported file: {file_path}")
                        continue
                    
                    # Generate document ID
                    doc_id = self.generate_doc_id(file_path)
                    
                    if last_sync_time is not None:
                        # Check if document has already been synced
                        if self.vector_store_manager and self.vector_store_manager.document_exists(doc_id + "_chunk_0"):
                            # If document is already synced, check if it's been modified since last sync
                            if last_sync_time:
                                file_stat = self._sftp_client.stat(file_path)
                                file_mtime = datetime.fromtimestamp(file_stat.st_mtime)
                                # Skip if not modified since last sync
                                if file_mtime <= last_sync_time:
                                    continue
                    
                    # Read file content
                    with self._sftp_client.open(file_path, 'rb') as f:
                        file_bytes = f.read()
                    
                    # Parse file content
                    if file_bytes:
                        try:
                            # 使用 FileParser 解析文件内容
                            parsed_docs = self.file_parser.parse_file_content(file_bytes, file_path, doc_id=doc_id, host=self.config.host)
                            if parsed_docs:
                                # 合并所有文档内容
                                content = '\n\n'.join(doc.text for doc in parsed_docs if doc.text)
                            else:
                                logger.warning(f"No content extracted from {file_path}")
                                content = f"[无法读取文件：{Path(file_path).name}]"
                        except Exception as e:
                            logger.error(f"Error parsing file content for {file_path}: {e}")
                            content = f"[无法读取文件：{Path(file_path).name}]"
                    else:
                        content = f"[无法读取文件：{Path(file_path).name}]"
                    
                    # Build document
                    document = {
                        'id': doc_id,
                        'content': content,
                        'metadata': {
                            'file_path': str(file_path),
                            'update_time': datetime.fromtimestamp(self._sftp_client.stat(file_path).st_mtime),
                            'host': self.config.host
                        }
                    }
                    documents.append(document)
                except Exception as e:
                    logger.error(f"Error processing file {file_path}: {e}")
        finally:
            self._disconnect()
        
        return documents
    
    def fetch_new_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
        """
        Fetch new/updated documents from the folder since last sync time
        
        Args:
            last_sync_time: Last synchronization time
            
        Returns:
            List of new/updated documents
        """
        # Call fetch_all_documents which now handles document existence checks
        return self.fetch_all_documents(last_sync_time)
    
    def get_synced_document_ids(self) -> Set[str]:
        """
        Get IDs of all files in the folder

        Returns:
            Set of file paths (as document IDs)
        """
        self._connect()
        try:
            files = self._get_all_files()
            return set(files)
        finally:
            self._disconnect()
    
    def generate_doc_id(self, file_path: str) -> str:
        """
        Generate a unique document ID for files
        
        Args:
            file_path: Path to the file
            
        Returns:
            Unique document ID based on host IP and file path
        """
        # 使用配置中的主机地址
        host_address = self.config.host or 'unknown'
        
        # 替换路径中的特殊字符，避免生成无效的doc_id
        sanitized_path = file_path.replace('/', '_').replace('\\', '_').replace(':', '_').replace(' ', '_')
        return f"{host_address}_{sanitized_path}"
    
    def doc_to_llamaindex_doc(self, doc: Dict) -> 'Document':
        """
        Convert folder document to LlamaIndex Document
        
        Args:
            doc: Folder document dictionary
            
        Returns:
            LlamaIndex Document object
        """
        content = doc.get('content', "")
        doc_id = doc.get('id', "")
        metadata = doc.get('metadata', {})
        
        # Ensure metadata has source information
        metadata['source'] = 'folder'
        metadata['host'] = self.config.host
        
        # Create Document
        return Document(
            text=content,
            id_=doc_id,
            metadata=metadata
        )
    
    def _connect(self):
        """
        Connect to the server via SSH/SFTP
        
        Raises:
            Exception: If connection fails with detailed error message
        """
        import paramiko
        
        self._ssh_client = paramiko.SSHClient()
        self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        # Connect to SSH server
        
        # 获取用户名
        username = self.config.username
        if not username:
            raise Exception("SSH connection failed: Username is required")
        
        # 建立SSH连接的参数
        ssh_params = {
            'hostname': self.config.host,
            'port': self.config.port or 22,
            'username': username,
            'password': self.config.password,
            'timeout': 10,
            'allow_agent': True,  # 允许使用SSH代理
            'look_for_keys': False  # 禁用查找本地密钥文件
        }
        
        try:
            # 连接到SSH服务器
            self._ssh_client.connect(**ssh_params)
            
            # Create SFTP client
            self._sftp_client = self._ssh_client.open_sftp()
        except paramiko.AuthenticationException:
            raise Exception(f"SSH connection failed: Authentication failed for user {username} on {self.config.host}")
        except paramiko.SSHException as ssh_error:
            raise Exception(f"SSH connection failed: {str(ssh_error)}")
        except Exception as e:
            raise Exception(f"Connection failed: {str(e)}")
    
    def _disconnect(self):
        """
        Disconnect from the server
        """
        if self._sftp_client:
            self._sftp_client.close()
            self._sftp_client = None
        
        if self._ssh_client:
            self._ssh_client.close()
            self._ssh_client = None
    
    def _get_all_files(self) -> List[str]:
        """
        Get all files in the folder
        
        Returns:
            List of file paths
            
        Note:
            This method assumes that a connection has already been established by the caller
        """
        files = []
        try:
            # 直接调用 _get_files_recursive，使用已经建立的连接
            self._get_files_recursive(self.config.folder_path, files)
        except Exception as e:
            logger.error(f"Error getting all files: {e}")
        return files
    
    def _get_files_recursive(self, folder_path: str, files: List[str]):
        """
        Recursively get all files in the folder
        
        Args:
            folder_path: Current folder path
            files: List to store found files
        """
        try:
            items = self._sftp_client.listdir_attr(folder_path)
            
            for item in items:
                item_path = os.path.join(folder_path, item.filename)
                
                if item.filename not in ('.', '..'):
                    if item.st_mode & 0o040000:  # Check if it's a directory
                        if self.config.recursive:
                            self._get_files_recursive(item_path, files)
                    else:
                        # Check if file should be ignored
                        if not self._should_ignore_file(item_path):
                            files.append(item_path)
        except Exception as e:
            logger.error(f"Error listing folder {folder_path}: {e}")
    
    def _should_ignore_file(self, file_path: str) -> bool:
        """
        Check if file should be ignored based on ignore patterns
        
        Args:
            file_path: File path to check
            
        Returns:
            True if file should be ignored, False otherwise
        """
        if not hasattr(self.config, 'ignore_patterns') or not self.config.ignore_patterns:
            return False
        
        # Get relative path from folder root
        relative_path = os.path.relpath(file_path, self.config.folder_path)
        
        for pattern in self.config.ignore_patterns:
            if self._match_pattern(relative_path, pattern):
                return True
        
        return False
    
    def _match_pattern(self, path: str, pattern: str) -> bool:
        """
        Match a path against a pattern (similar to .gitignore)
        
        Args:
            path: Path to match
            pattern: Pattern to match against
            
        Returns:
            True if path matches pattern, False otherwise
        """
        # Convert glob pattern to regex
        regex_pattern = pattern
        regex_pattern = regex_pattern.replace('.', r'\.')
        regex_pattern = regex_pattern.replace('*', r'.*')
        regex_pattern = regex_pattern.replace('?', r'.')
        
        # Handle directory patterns
        if pattern.endswith('/'):
            regex_pattern = f'^{regex_pattern}.*$'
        else:
            regex_pattern = f'^{regex_pattern}$'
        
        return bool(re.match(regex_pattern, path))
    
    @staticmethod
    def check_data_source_exists(config: BaseDataSourceConfig) -> bool:
        """
        Check if the folder exists and is accessible
        
        Args:
            config: Folder configuration
            
        Returns:
            True if folder exists and is accessible, False otherwise
            
        Raises:
            Exception: If connection fails with detailed error message
        """
        import paramiko
        
        ssh_client = None
        sftp_client = None
        try:
            # 检查必要的配置
            if not config.host:
                raise Exception("SSH connection failed: Host is required")
            
            username = config.username
            if not username:
                raise Exception("SSH connection failed: Username is required")
            
            ssh_client = paramiko.SSHClient()
            ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # Connect to SSH server - use SSH agent if available, otherwise password
            try:
                ssh_client.connect(
                    hostname=config.host,
                    port=config.port or 22,
                    username=username,
                    password=config.password,
                    timeout=10,
                    allow_agent=True,  # 允许使用SSH代理
                    look_for_keys=False  # 禁用查找本地密钥文件
                )
            except paramiko.AuthenticationException:
                raise Exception(f"SSH connection failed: Authentication failed for user {username} on {config.host}")
            except paramiko.SSHException as ssh_error:
                raise Exception(f"SSH connection failed: {str(ssh_error)}")
            except Exception as e:
                raise Exception(f"Connection failed: {str(e)}")
            
            # Create SFTP client and check folder exists
            try:
                sftp_client = ssh_client.open_sftp()
                sftp_client.stat(config.folder_path)
            except Exception as e:
                raise Exception(f"Folder access failed: {str(e)}")
            
            return True
        except Exception as e:
            logger.error(f"Error checking folder: {e}")
            # 重新抛出异常，以便上层能够捕获并传递详细的错误信息
            raise
        finally:
            if sftp_client:
                sftp_client.close()
            if ssh_client:
                ssh_client.close()