299 lines
10 KiB
Python
299 lines
10 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
|
|
logger.error(f"Error getting file bytes 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 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
|
|
documents = []
|
|
for doc in docs:
|
|
try:
|
|
llamaindex_doc = self.doc_to_llamaindex_doc(doc)
|
|
if len(llamaindex_doc.text.strip()) >= 1:
|
|
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}")
|
|
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,
|
|
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'))
|
|
logger.error(f"Error chunking document {doc_id}: {e}")
|
|
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:
|
|
# 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, 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.folder_sync import FolderSync
|
|
|
|
sync_classes = {
|
|
'database': MySQLSync, # Currently only MySQL, but can be extended
|
|
'folder': FolderSync
|
|
}
|
|
|
|
if source_type not in sync_classes:
|
|
raise ValueError(f"Unsupported data source type: {source_type}")
|
|
|
|
return sync_classes[source_type]
|