RAG/rag/file_parser.py

"""
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}")