444 lines
16 KiB
Python
444 lines
16 KiB
Python
"""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:
|
||
# Check if it's a local test server
|
||
if self.host in ['localhost', '127.0.0.1']:
|
||
# For local testing, directly return success
|
||
# because we've already mounted the directory via volume
|
||
logger.info(f"Local SSH connection test successful: {self.host}")
|
||
return True
|
||
|
||
# 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()
|