203 lines
8.8 KiB
Python
203 lines
8.8 KiB
Python
"""
|
||
Document processing and chunking module
|
||
"""
|
||
from llama_index.core import Document
|
||
from llama_index.core.node_parser import SentenceSplitter
|
||
from typing import List, Dict, Optional
|
||
from datetime import datetime, date
|
||
from loguru import logger
|
||
from config import settings
|
||
|
||
|
||
class DocumentProcessor:
|
||
"""Process and chunk documents for RAG"""
|
||
|
||
def __init__(self):
|
||
self.node_parser = SentenceSplitter(
|
||
chunk_size=settings.CHUNK_SIZE,
|
||
chunk_overlap=settings.CHUNK_OVERLAP
|
||
)
|
||
|
||
def mysql_doc_to_llamaindex_doc(self, mysql_doc: Dict, db_config=None) -> Document:
|
||
"""
|
||
Convert MySQL document to LlamaIndex Document
|
||
将单个MySQL文档转换为LlamaIndex文档对象
|
||
|
||
Args:
|
||
mysql_doc: Dictionary from MySQL query result
|
||
db_config: DatabaseConfig object (optional, for backward compatibility)
|
||
|
||
Returns:
|
||
LlamaIndex Document object
|
||
"""
|
||
# Get column names from db_config or fallback to default values
|
||
id_column = db_config.id_column if db_config else "id"
|
||
title_column = db_config.title_column if db_config else "title"
|
||
metadata_columns = db_config.metadata_columns if db_config else None
|
||
|
||
doc_id = str(mysql_doc.get(id_column, ""))
|
||
|
||
# 处理多个 content 列（支持合并多个列的内容）
|
||
if db_config:
|
||
# 使用配置的多个 content 列
|
||
content_columns = db_config.content_columns
|
||
content_separator = db_config.content_separator
|
||
else:
|
||
# 向后兼容：使用单个 content_column
|
||
content_columns = ["content"] # Default to "content" column
|
||
content_separator = "\n"
|
||
|
||
# 合并所有 content 列的内容
|
||
content_parts = []
|
||
for col in content_columns:
|
||
col_value = mysql_doc.get(col, "")
|
||
if col_value:
|
||
content_parts.append(str(col_value))
|
||
|
||
# 用指定的分隔符连接多个列的内容
|
||
content = content_separator.join(content_parts) if content_parts else ""
|
||
|
||
title = mysql_doc.get(title_column, "") if title_column else None
|
||
|
||
# Create unique doc_id with database source and table name to avoid conflicts
|
||
# Format: {db_source}_{table_name}_{id} to ensure uniqueness across multiple tables
|
||
# This must match the format used in sync_service.py
|
||
if '_db_source' in mysql_doc and '_db_table' in mysql_doc:
|
||
unique_doc_id = f"{mysql_doc['_db_source']}_{mysql_doc['_db_table']}_{doc_id}"
|
||
elif '_db_source' in mysql_doc:
|
||
unique_doc_id = f"{mysql_doc['_db_source']}_{doc_id}"
|
||
else:
|
||
unique_doc_id = doc_id
|
||
|
||
# Build metadata
|
||
# Note: metadata['doc_id'] should use unique_doc_id to match sync_service.py logic
|
||
# Format: {db_source}_{table_name}_{id} to ensure uniqueness across multiple tables
|
||
metadata = {
|
||
"doc_id": unique_doc_id, # Format: {db_source}_{table_name}_{id}
|
||
"original_doc_id": doc_id, # Keep original ID for reference
|
||
"source": "mysql",
|
||
"content_column": db_config.content_column # 额外存入 content_column
|
||
}
|
||
|
||
# Add database source information if available
|
||
if '_db_source' in mysql_doc:
|
||
metadata['db_source'] = mysql_doc['_db_source']
|
||
if '_db_database' in mysql_doc:
|
||
metadata['db_database'] = mysql_doc['_db_database']
|
||
if '_db_table' in mysql_doc:
|
||
metadata['db_table'] = mysql_doc['_db_table']
|
||
|
||
if title:
|
||
metadata["title"] = title
|
||
|
||
# Add additional metadata columns if specified
|
||
if metadata_columns:
|
||
metadata_cols = [col.strip() for col in metadata_columns.split(",")]
|
||
for col in metadata_cols:
|
||
if col in mysql_doc and not col.startswith('_'):
|
||
value = mysql_doc[col]
|
||
# Convert datetime/date objects to strings for ChromaDB compatibility
|
||
# ChromaDB metadata only supports: str, int, float, bool
|
||
if isinstance(value, (datetime, date)):
|
||
# Format datetime/date as ISO format string
|
||
metadata[col] = value.isoformat()
|
||
elif value is not None:
|
||
# Convert other types to string if not already a supported type
|
||
if not isinstance(value, (str, int, float, bool)):
|
||
metadata[col] = str(value)
|
||
else:
|
||
metadata[col] = value
|
||
|
||
# Create Document
|
||
# Note: LlamaIndex uses id_ (not doc_id) as the unique identifier for Document
|
||
# See: https://docs.llamaindex.org.cn/en/stable/module_guides/indexing/document_management/#update
|
||
doc = Document(
|
||
text=content,
|
||
id_=unique_doc_id, # Use id_ as per LlamaIndex documentation
|
||
metadata=metadata
|
||
)
|
||
|
||
return doc
|
||
|
||
def process_documents(self, mysql_docs: List[Dict], db_config=None) -> List[Document]:
|
||
"""
|
||
Process multiple MySQL documents into LlamaIndex Documents
|
||
处理多个MySQL文档，将其转换为LlamaIndex文档对象
|
||
Args:
|
||
mysql_docs: List of MySQL document dictionaries
|
||
db_config: DatabaseConfig object (optional, for backward compatibility)
|
||
|
||
Returns:
|
||
List of LlamaIndex Document objects
|
||
"""
|
||
documents = []
|
||
id_column = db_config.id_column if db_config else "id"
|
||
for mysql_doc in mysql_docs:
|
||
try:
|
||
doc = self.mysql_doc_to_llamaindex_doc(mysql_doc, db_config)
|
||
if len(doc.text.strip()) >= 100:
|
||
documents.append(doc)
|
||
else:
|
||
logger.warning(f"跳过过短文档 (id: {doc.id_})，内容长度: {len(doc.text)} 字符")
|
||
except Exception as e:
|
||
doc_id = mysql_doc.get(id_column, mysql_doc.get('_db_source', 'unknown'))
|
||
logger.error(f"Error processing document {doc_id}: {e}")
|
||
continue
|
||
|
||
logger.info(f"Processed {len(documents)} documents")
|
||
return documents
|
||
|
||
def chunk_documents(self, documents: List[Document]) -> List[Document]:
|
||
"""
|
||
Chunk documents into smaller pieces
|
||
|
||
Args:
|
||
documents: List of Document objects
|
||
doc = Document(
|
||
text=content,
|
||
id_=unique_doc_id, # Use id_ as per LlamaIndex documentation
|
||
metadata=metadata
|
||
)
|
||
|
||
Returns:
|
||
List of chunked Document objects
|
||
"""
|
||
chunked_docs = []
|
||
for doc in documents:
|
||
try:
|
||
nodes = self.node_parser.get_nodes_from_documents([doc])
|
||
# Get unique_doc_id from doc.id_ or metadata['doc_id']
|
||
# This is the original document's unique ID (format: {db_source}_{table_name}_{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
|
||
# Use unique_doc_id + chunk_index as id_ for easier identification and duplicate checking
|
||
# Format: {db_source}_{table_name}_{id}_chunk_{index}
|
||
for index, node in enumerate(nodes):
|
||
chunk_id = f"{unique_doc_id}_chunk_{index}"
|
||
chunked_doc = Document(
|
||
text=node.text,
|
||
id_=chunk_id, # Use unique_doc_id + chunk_index for better traceability
|
||
metadata={
|
||
**doc.metadata,
|
||
"chunk_id": chunk_id,
|
||
"chunk_index": index,
|
||
"total_chunks": len(nodes),
|
||
"original_node_id": node.node_id # Keep original node_id for reference
|
||
}
|
||
)
|
||
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
|
||
|