249 lines
9.5 KiB
Python
249 lines
9.5 KiB
Python
"""
|
||
File parsing module for various document formats
|
||
"""
|
||
import os
|
||
import mimetypes
|
||
from typing import List, Dict, Optional
|
||
from pathlib import Path
|
||
from loguru import logger
|
||
from llama_index.core import Document
|
||
from llama_index.readers.alibabacloud_aisearch import (
|
||
AlibabaCloudAISearchDocumentReader,
|
||
AlibabaCloudAISearchImageReader,
|
||
)
|
||
from llama_index.core.readers import SimpleDirectoryReader
|
||
|
||
|
||
class FileParser:
|
||
"""Parse various file formats into LlamaIndex Documents"""
|
||
|
||
# Supported file extensions
|
||
SUPPORTED_EXTENSIONS = {
|
||
'.txt', '.md', '.markdown', # Text files
|
||
'.pdf', # PDF files
|
||
'.docx', '.doc', # Word documents
|
||
'.html', '.htm', # HTML files
|
||
'.csv', # CSV files
|
||
'.json', # JSON files
|
||
}
|
||
|
||
def __init__(self):
|
||
"""Initialize file parser"""
|
||
pass
|
||
|
||
def is_supported(self, filename: str) -> bool:
|
||
"""
|
||
Check if file format is supported
|
||
|
||
Args:
|
||
filename: File name or path
|
||
|
||
Returns:
|
||
True if supported, False otherwise
|
||
"""
|
||
ext = Path(filename).suffix.lower() # 获取文件类型并转为小写
|
||
return ext in self.SUPPORTED_EXTENSIONS
|
||
|
||
def parse_file(self, file_path: str, doc_id: Optional[str] = None, metadata: Optional[Dict] = None) -> List[Document]:
|
||
"""
|
||
Parse a file into LlamaIndex Documents
|
||
|
||
Args:
|
||
file_path: Path to the file
|
||
doc_id: Optional document ID (if not provided, will use filename)
|
||
metadata: Optional metadata to add to documents
|
||
|
||
Returns:
|
||
List of LlamaIndex Document objects
|
||
"""
|
||
if not os.path.exists(file_path):
|
||
raise FileNotFoundError(f"File not found: {file_path}")
|
||
|
||
filename = os.path.basename(file_path)
|
||
ext = Path(file_path).suffix.lower()
|
||
|
||
if not self.is_supported(file_path):
|
||
raise ValueError(f"Unsupported file format: {ext}. Supported formats: {', '.join(self.SUPPORTED_EXTENSIONS)}")
|
||
|
||
# Use doc_id or generate from filename
|
||
if not doc_id:
|
||
doc_id = Path(filename).stem
|
||
|
||
try:
|
||
logger.info(f"Starting to parse file: {filename} (type: {ext})")
|
||
|
||
# 阿里巴巴的llama index reader插件可以解析doc文档
|
||
document_reader = AlibabaCloudAISearchDocumentReader()
|
||
image_reader = AlibabaCloudAISearchImageReader()
|
||
|
||
file_extractor = {}
|
||
for suffix in (".pdf", ".docx", ".doc", ".ppt", ".pptx"):
|
||
file_extractor[suffix] = document_reader
|
||
for suffix in (".jpg", ".jpeg", ".png", ".bmp", ".tiff"):
|
||
file_extractor[suffix] = image_reader
|
||
|
||
# Use LlamaIndex's SimpleDirectoryReader for parsing
|
||
# It supports many formats out of the box
|
||
reader = SimpleDirectoryReader(
|
||
input_files=[file_path],
|
||
file_extractor=file_extractor,
|
||
filename_as_id=False # We'll set custom IDs
|
||
)
|
||
|
||
# Read documents with timeout protection
|
||
logger.debug(f"Loading data from file: {file_path}")
|
||
documents = reader.load_data()
|
||
logger.debug(f"Loaded {len(documents)} document(s) from file")
|
||
|
||
# Validate that we got documents
|
||
if not documents:
|
||
logger.warning(f"No documents extracted from file: {filename}")
|
||
raise ValueError(f"Failed to extract content from {filename}. The file may be empty or in an unsupported format.")
|
||
|
||
# Check if documents have content
|
||
empty_docs = []
|
||
for i, doc in enumerate(documents):
|
||
if not doc.text or not doc.text.strip():
|
||
empty_docs.append(i)
|
||
logger.warning(f"Document {i} from {filename} has empty content")
|
||
|
||
if len(empty_docs) == len(documents):
|
||
raise ValueError(f"All documents from {filename} have empty content. The file may not contain readable text.")
|
||
|
||
# Add metadata to each document
|
||
base_metadata = {
|
||
'source': 'file_upload',
|
||
'file_name': filename,
|
||
'file_path': file_path,
|
||
'file_type': ext.lstrip('.'),
|
||
}
|
||
|
||
if metadata:
|
||
base_metadata.update(metadata)
|
||
|
||
# Update documents with metadata and doc_id
|
||
valid_documents = []
|
||
for i, doc in enumerate(documents):
|
||
# Skip empty documents
|
||
if not doc.text or not doc.text.strip():
|
||
logger.warning(f"Skipping empty document {i} from {filename}")
|
||
continue
|
||
|
||
# Set document ID
|
||
if len(documents) == 1:
|
||
# Single document, use doc_id directly
|
||
doc.id_ = doc_id
|
||
doc.metadata = {**base_metadata, 'doc_id': doc_id}
|
||
else:
|
||
# Multiple documents (e.g., PDF pages), append index
|
||
chunk_id = f"{doc_id}_chunk_{i}"
|
||
doc.id_ = chunk_id
|
||
doc.metadata = {**base_metadata, 'doc_id': doc_id, 'chunk_id': chunk_id, 'chunk_index': i}
|
||
|
||
valid_documents.append(doc)
|
||
|
||
if not valid_documents:
|
||
raise ValueError(f"No valid documents extracted from {filename}. All documents are empty.")
|
||
|
||
logger.info(f"Successfully parsed file {filename}: {len(valid_documents)} valid document(s) (total: {len(documents)}, skipped empty: {len(documents) - len(valid_documents)})")
|
||
return valid_documents
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error parsing file {file_path}: {e}", exc_info=True)
|
||
raise
|
||
|
||
def parse_file_content(self, content: bytes, filename: str, doc_id: Optional[str] = None, metadata: Optional[Dict] = None) -> List[Document]:
|
||
"""
|
||
Parse file content from bytes into LlamaIndex Documents
|
||
|
||
Args:
|
||
content: File content as bytes
|
||
filename: Original filename (for format detection)
|
||
doc_id: Optional document ID
|
||
metadata: Optional metadata to add to documents
|
||
|
||
Returns:
|
||
List of LlamaIndex Document objects
|
||
"""
|
||
import tempfile
|
||
|
||
# Create temporary file
|
||
ext = Path(filename).suffix.lower() # 确定文件扩展名
|
||
if not ext:
|
||
# 无拓展名时，Try to detect from mimetype
|
||
mime_type, _ = mimetypes.guess_type(filename)
|
||
if mime_type:
|
||
ext = mimetypes.guess_extension(mime_type) or '.txt'
|
||
else:
|
||
ext = '.txt'
|
||
|
||
# Use tempfile to save content and parse
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
|
||
tmp_file.write(content)
|
||
tmp_path = tmp_file.name
|
||
|
||
try:
|
||
# Parse the temporary file
|
||
documents = self.parse_file(tmp_path, doc_id=doc_id, metadata=metadata)
|
||
return documents
|
||
finally:
|
||
# Clean up temporary file
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to delete temporary file {tmp_path}: {e}")
|
||
|
||
|
||
def get_text_from_bytes(self, file_bytes: bytes, file_name: str) -> str:
|
||
"""
|
||
根据文件的字节流获取文件内的文本字符串
|
||
Args:
|
||
file_bytes: 文件字节流
|
||
file_name: 文件名
|
||
Returns:
|
||
文件内容的文本字符串
|
||
"""
|
||
import tempfile
|
||
|
||
# 获取文件类型
|
||
if file_name == "unknown":
|
||
ext = '.txt'
|
||
else:
|
||
ext = Path(file_name).suffix.lower()
|
||
|
||
# Use tempfile to save content and parse
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
|
||
tmp_file.write(file_bytes)
|
||
tmp_path = tmp_file.name
|
||
|
||
try:
|
||
# Parse the temporary file
|
||
documents = self.parse_file(tmp_path, doc_id=None)
|
||
|
||
# 提取所有非空文档的文本内容
|
||
text_parts = []
|
||
for i, doc in enumerate(documents):
|
||
if doc.text and doc.text.strip():
|
||
text_parts.append(doc.text.strip())
|
||
logger.debug(f"提取到文档 {i} 的文本内容，长度: {len(doc.text)}")
|
||
else:
|
||
logger.warning(f"文档 {i} 内容为空，已跳过")
|
||
|
||
if not text_parts:
|
||
raise ValueError("所有文档内容均为空，无法提取有效文本")
|
||
|
||
# 合并所有文本内容
|
||
full_text = "\n\n".join(text_parts)
|
||
logger.info(f"成功从文件(文件名：{file_name})字节流中提取文本，总长度: {len(full_text)} 字符")
|
||
return full_text
|
||
|
||
except Exception as e:
|
||
logger.error(f"解析文件(文件名：{file_name})字节流时出错: {e}", exc_info=True)
|
||
raise
|
||
finally:
|
||
# Clean up temporary file
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to delete temporary file {tmp_path}: {e}")
|