168 lines
6.8 KiB
Python
168 lines
6.8 KiB
Python
"""
|
||
File parsing module for various document formats
|
||
"""
|
||
import os
|
||
import requests
|
||
from config import settings
|
||
import mimetypes
|
||
from typing import List, Dict, Optional
|
||
from pathlib import Path
|
||
from loguru import logger
|
||
from config import settings
|
||
from llama_index.core import Document
|
||
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 the file format is supported by the parser
|
||
|
||
Args:
|
||
filename: Name or path of the file to check
|
||
|
||
Returns:
|
||
True if supported, False otherwise
|
||
"""
|
||
ext = Path(filename).suffix.lower() # 获取文件类型并转为小写
|
||
return ext in self.SUPPORTED_EXTENSIONS
|
||
|
||
def parse_file_content(self, content: bytes, file_path: str, doc_id: Optional[str] = None, metadata: Optional[Dict] = None, host: Optional[str] = None) -> List[Document]:
|
||
"""
|
||
Parse file content from bytes into LlamaIndex Documents
|
||
|
||
Args:
|
||
content: File content as bytes
|
||
file_path: Original file path (for format detection and metadata)
|
||
doc_id: Optional document ID
|
||
metadata: Optional metadata to add to documents
|
||
host: Optional host address (for remote files, use the remote host address)
|
||
|
||
Returns:
|
||
List of LlamaIndex Document objects
|
||
"""
|
||
import tempfile
|
||
|
||
# Extract filename from path
|
||
original_filename = Path(file_path).name
|
||
ext = Path(file_path).suffix.lower() # 确定文件扩展名
|
||
|
||
if not ext:
|
||
# 无拓展名时，Try to detect from mimetype
|
||
mime_type, _ = mimetypes.guess_type(file_path)
|
||
if mime_type:
|
||
ext = mimetypes.guess_extension(mime_type) or '.txt'
|
||
else:
|
||
ext = '.txt'
|
||
|
||
# Check if file format is supported
|
||
if not self.is_supported(file_path):
|
||
raise ValueError(f"Unsupported file format: {ext}. Supported formats: {', '.join(self.SUPPORTED_EXTENSIONS)}")
|
||
|
||
try:
|
||
logger.info(f"Starting to parse file: {original_filename} (type: {ext})\n")
|
||
|
||
# 处理 .doc 文件：通过 soffice-service 转换为 .docx
|
||
if ext == '.doc':
|
||
logger.info(f"检测到 .doc 文件，开始转换为 .docx 格式: {file_path}")
|
||
# 上传文件到 soffice-service 的 convert 接口
|
||
files = {'file': (original_filename, content, 'application/msword')}
|
||
soffice_url = f"http://{settings.SOFFICE_HOST}:{settings.SOFFICE_PORT}/convert"
|
||
response = requests.post(soffice_url, files=files, timeout=60)
|
||
try:
|
||
response.raise_for_status() # 检查 HTTP 错误
|
||
except requests.exceptions.HTTPError as e:
|
||
logger.error(f"转换请求失败: {e}")
|
||
raise
|
||
# 将转换后的 docx 内容保存到临时文件进行解析
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as tmp_docx_file:
|
||
tmp_docx_file.write(response.content)
|
||
parse_path = tmp_docx_file.name
|
||
else:
|
||
# 对于其他文件格式，创建临时文件保存内容
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
|
||
tmp_file.write(content)
|
||
parse_path = tmp_file.name
|
||
|
||
# Use LlamaIndex's SimpleDirectoryReader for parsing
|
||
# It supports many formats out of the box
|
||
reader = SimpleDirectoryReader(
|
||
input_files=[parse_path],
|
||
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: {original_filename}")
|
||
raise ValueError(f"Failed to extract content from {original_filename}. The file may be empty or in an unsupported format.")
|
||
|
||
# Add metadata to each document
|
||
base_metadata = {
|
||
'source': 'file_upload',
|
||
'file_type': ext.lstrip('.'),
|
||
'file_name': original_filename,
|
||
'file_path': file_path, # Use the original file path
|
||
}
|
||
|
||
# Update with provided metadata if any
|
||
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 {original_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)
|
||
|
||
logger.info(f"Successfully parsed file {original_filename}: {len(valid_documents)} valid document(s) (total: {len(documents)}, skipped empty: {len(documents) - len(valid_documents)})\n")
|
||
return valid_documents
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error parsing file {file_path}: {e}", exc_info=True)
|
||
raise
|
||
finally:
|
||
try:
|
||
os.unlink(parse_path)
|
||
logger.debug(f"Cleaned up temporary file: {parse_path}")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to delete temporary file {parse_path}: {e}")
|
||
|
||
|
||
|