315 lines
11 KiB
Python
315 lines
11 KiB
Python
"""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}")
|