DocsGPT/application/parser/file/bulk.py

"""Simple reader that reads files of different formats from a directory."""
import logging
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple, Union

from application.parser.file.base import BaseReader
from application.parser.file.base_parser import BaseParser, DocumentParseError
from application.parser.file.docs_parser import DocxParser, PDFParser
from application.parser.file.epub_parser import EpubParser
from application.parser.file.html_parser import HTMLParser
from application.parser.file.markdown_parser import MarkdownParser
from application.parser.file.rst_parser import RstParser
from application.parser.file.tabular_parser import PandasCSVParser, ExcelParser
from application.parser.file.json_parser import JSONParser
from application.parser.file.pptx_parser import PPTXParser
from application.parser.file.image_parser import ImageParser
from application.parser.file.audio_parser import AudioParser
from application.parser.schema.base import Document
from application.stt.constants import SUPPORTED_AUDIO_EXTENSIONS
from application.utils import num_tokens_from_string
from application.core.settings import settings


def _build_audio_parser_mapping() -> Dict[str, BaseParser]:
    return {extension: AudioParser() for extension in SUPPORTED_AUDIO_EXTENSIONS}


def get_default_file_extractor(
    ocr_enabled: Optional[bool] = None,
) -> Dict[str, BaseParser]:
    """Get the default file extractor.

    Uses docling parsers by default for advanced document processing.
    Falls back to standard parsers if docling is not installed.
    """
    try:
        from application.parser.file.docling_parser import (
            DoclingPDFParser,
            DoclingDocxParser,
            DoclingPPTXParser,
            DoclingXLSXParser,
            DoclingHTMLParser,
            DoclingImageParser,
            DoclingCSVParser,
            DoclingAsciiDocParser,
            DoclingVTTParser,
            DoclingXMLParser,
        )
        if ocr_enabled is None:
            ocr_enabled = settings.DOCLING_OCR_ENABLED
        return {
            # Documents
            ".pdf": DoclingPDFParser(ocr_enabled=ocr_enabled),
            ".docx": DoclingDocxParser(),
            ".pptx": DoclingPPTXParser(),
            ".xlsx": DoclingXLSXParser(),
            # Web formats
            ".html": DoclingHTMLParser(),
            ".xhtml": DoclingHTMLParser(),
            # Data formats
            ".csv": DoclingCSVParser(),
            ".json": JSONParser(),  # Keep JSON parser (specialized handling)
            # Text/markup formats
            ".md": MarkdownParser(),  # Keep markdown parser (specialized handling)
            ".mdx": MarkdownParser(),
            ".rst": RstParser(),
            ".adoc": DoclingAsciiDocParser(),
            ".asciidoc": DoclingAsciiDocParser(),
            # Images (with OCR) - only use Docling when OCR is enabled
            ".png": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            ".jpg": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            ".jpeg": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            ".tiff": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            ".tif": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            ".bmp": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            ".webp": DoclingImageParser(ocr_enabled=ocr_enabled) if ocr_enabled else ImageParser(),
            # Media/subtitles
            ".vtt": DoclingVTTParser(),
            **_build_audio_parser_mapping(),
            # Specialized XML formats
            ".xml": DoclingXMLParser(),
            # Formats docling doesn't support - use standard parsers
            ".epub": EpubParser(),
        }
    except ImportError:
        logging.warning(
            "docling is not installed. Using standard parsers. "
            "For advanced document parsing, install with: pip install docling"
        )
        # Fallback to standard parsers
        return {
            ".pdf": PDFParser(),
            ".docx": DocxParser(),
            ".csv": PandasCSVParser(),
            ".xlsx": ExcelParser(),
            ".epub": EpubParser(),
            ".md": MarkdownParser(),
            ".rst": RstParser(),
            ".html": HTMLParser(),
            ".mdx": MarkdownParser(),
            ".json": JSONParser(),
            ".pptx": PPTXParser(),
            ".png": ImageParser(),
            ".jpg": ImageParser(),
            ".jpeg": ImageParser(),
            **_build_audio_parser_mapping(),
        }


# For backwards compatibility
DEFAULT_FILE_EXTRACTOR: Dict[str, BaseParser] = get_default_file_extractor()


class SimpleDirectoryReader(BaseReader):
    """Simple directory reader.

    Can read files into separate documents, or concatenates
    files into one document text.

    Args:
        input_dir (str): Path to the directory.
        input_files (List): List of file paths to read (Optional; overrides input_dir)
        exclude_hidden (bool): Whether to exclude hidden files (dotfiles).
        errors (str): how encoding and decoding errors are to be handled,
              see https://docs.python.org/3/library/functions.html#open
        recursive (bool): Whether to recursively search in subdirectories.
            False by default.
        required_exts (Optional[List[str]]): List of required extensions.
            Default is None.
        file_extractor (Optional[Dict[str, BaseParser]]): A mapping of file
            extension to a BaseParser class that specifies how to convert that file
            to text. See DEFAULT_FILE_EXTRACTOR.
        num_files_limit (Optional[int]): Maximum number of files to read.
            Default is None.
        file_metadata (Optional[Callable[str, Dict]]): A function that takes
            in a filename and returns a Dict of metadata for the Document.
            Default is None.
    """

    def __init__(
            self,
            input_dir: Optional[str] = None,
            input_files: Optional[List] = None,
            exclude_hidden: bool = True,
            errors: str = "ignore",
            recursive: bool = True,
            required_exts: Optional[List[str]] = None,
            file_extractor: Optional[Dict[str, BaseParser]] = None,
            num_files_limit: Optional[int] = None,
            file_metadata: Optional[Callable[[str], Dict]] = None,
    ) -> None:
        """Initialize with parameters."""
        super().__init__()

        if not input_dir and not input_files:
            raise ValueError("Must provide either `input_dir` or `input_files`.")

        self.errors = errors

        self.recursive = recursive
        self.exclude_hidden = exclude_hidden
        # Normalize extensions to lowercase for case-insensitive matching
        self.required_exts = (
            [ext.lower() for ext in required_exts] if required_exts else None
        )
        self.num_files_limit = num_files_limit

        if input_files:
            self.input_files = []
            for path in input_files:
                print(path)
                input_file = Path(path)
                self.input_files.append(input_file)
        elif input_dir:
            self.input_dir = Path(input_dir)
            self.input_files = self._add_files(self.input_dir)

        self.file_extractor = file_extractor or DEFAULT_FILE_EXTRACTOR
        self.file_metadata = file_metadata
        # (path, message) per file skipped by ``load_data`` as unparseable.
        self.failed_files: List[Tuple[Path, str]] = []

    def _add_files(self, input_dir: Path) -> List[Path]:
        """Add files."""
        input_files = sorted(input_dir.iterdir())
        new_input_files = []
        dirs_to_explore = []
        for input_file in input_files:
            if input_file.is_dir():
                if self.recursive:
                    dirs_to_explore.append(input_file)
            elif self.exclude_hidden and input_file.name.startswith("."):
                continue
            elif (
                    self.required_exts is not None
                    and input_file.suffix.lower() not in self.required_exts
            ):
                continue
            else:
                new_input_files.append(input_file)

        for dir_to_explore in dirs_to_explore:
            sub_input_files = self._add_files(dir_to_explore)
            new_input_files.extend(sub_input_files)

        if self.num_files_limit is not None and self.num_files_limit > 0:
            new_input_files = new_input_files[0: self.num_files_limit]

        # print total number of files added
        logging.debug(
            f"> [SimpleDirectoryReader] Total files added: {len(new_input_files)}"
        )

        return new_input_files

    def load_data(
        self,
        concatenate: bool = False,
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> List[Document]:
        """Load data from the input directory.

        Args:
            concatenate (bool): whether to concatenate all files into one document.
                If set to True, file metadata is ignored.
                False by default.
            progress_callback (Optional[Callable[[int, int], None]]): Called
                after each file is parsed with ``(files_done, total_files)``.
                Lets callers surface parse/OCR progress before embedding
                begins. Exceptions raised by the callback are swallowed so
                progress reporting can never fail ingestion.

        Returns:
            List[Document]: A list of documents.

        Raises:
            DocumentParseError: if no input file could be parsed. Individual
                unreadable files are skipped and recorded in ``failed_files``
                so one corrupt document does not cost the caller the rest of
                the batch; a single-file read (the attachment path) still
                raises, since skipping there would only defer the failure to
                an empty result.
        """
        data: Union[str, List[str]] = ""
        data_list: List[str] = []
        metadata_list = []
        self.file_token_counts = {}
        self.failed_files = []

        total_files = len(self.input_files)

        def report_progress(files_done: int) -> None:
            if progress_callback is None:
                return
            try:
                progress_callback(files_done, total_files)
            except Exception:
                logging.warning("load_data progress callback failed", exc_info=True)

        for file_index, input_file in enumerate(self.input_files):
            suffix_lower = input_file.suffix.lower()
            parser_metadata = {}
            try:
                if suffix_lower in self.file_extractor:
                    parser = self.file_extractor[suffix_lower]
                    if not parser.parser_config_set:
                        parser.init_parser()
                    data = parser.parse_file(input_file, errors=self.errors)
                    parser_metadata = parser.get_file_metadata(input_file)
                else:
                    # do standard read
                    with open(input_file, "r", errors=self.errors) as f:
                        data = f.read()
            except DocumentParseError as e:
                logging.warning(f"Skipping unreadable file {input_file.name}: {e}")
                self.failed_files.append((input_file, str(e)))
                report_progress(file_index + 1)
                continue

            # Calculate token count for this file
            if isinstance(data, List):
                file_tokens = sum(num_tokens_from_string(str(d)) for d in data)
            else:
                file_tokens = num_tokens_from_string(str(data))
            
            full_path = str(input_file.resolve())
            self.file_token_counts[full_path] = file_tokens
            
            base_metadata = {
                'title': input_file.name,
                'token_count': file_tokens,
            }
            if parser_metadata:
                base_metadata.update(parser_metadata)
            
            if hasattr(self, 'input_dir'):
                try:
                    relative_path = str(input_file.relative_to(self.input_dir))
                    base_metadata['source'] = relative_path
                except ValueError:
                    base_metadata['source'] = str(input_file)
            else:
                base_metadata['source'] = str(input_file)

            if self.file_metadata is not None:
                custom_metadata = self.file_metadata(input_file.name)
                base_metadata.update(custom_metadata)

            if isinstance(data, List):
                # Extend data_list with each item in the data list
                data_list.extend([str(d) for d in data])
                # copy(): chunking writes token_count into this dict in
                # place, so a shared reference gives every chunk the last
                # chunk's count.
                metadata_list.extend([base_metadata.copy() for _ in data])
            else:
                data_list.append(str(data))
                metadata_list.append(base_metadata.copy())

            report_progress(file_index + 1)

        # Every file failed: there is nothing to ingest, so this is a failed
        # read rather than an empty one. Callers (the attachment worker, the
        # ingest tasks) treat it as terminal and tell the user.
        if self.failed_files and not data_list:
            if len(self.failed_files) == 1:
                # Single-file read (the attachment path): the parser's own
                # message reaches the user verbatim, so don't wrap it in
                # "None of the 1 file(s)…".
                raise DocumentParseError(self.failed_files[0][1])
            names = ", ".join(p.name for p, _ in self.failed_files[:5])
            raise DocumentParseError(
                f"None of the {len(self.failed_files)} files could be parsed: "
                f"{names}{'…' if len(self.failed_files) > 5 else ''} "
                f"({self.failed_files[0][1]})"
            )

        # Build directory structure if input_dir is provided
        if hasattr(self, 'input_dir'):
            self.directory_structure = self.build_directory_structure(self.input_dir)
            logging.info("Directory structure built successfully")
        else:
            self.directory_structure = {}

        if concatenate:
            return [Document("\n".join(data_list))]
        elif self.file_metadata is not None:
            return [Document(d, extra_info=m) for d, m in zip(data_list, metadata_list)]
        else:
            return [Document(d) for d in data_list]

    def build_directory_structure(self, base_path):
        """Build a dictionary representing the directory structure.

        Args:
            base_path: The base path to start building the structure from.

        Returns:
            dict: A nested dictionary representing the directory structure.
        """
        import mimetypes
        
        def build_tree(path):
            """Helper function to recursively build the directory tree."""
            result = {}
            
            for item in path.iterdir():
                if self.exclude_hidden and item.name.startswith('.'):
                    continue
                    
                if item.is_dir():
                    subtree = build_tree(item)
                    if subtree:
                        result[item.name] = subtree
                else:
                    if self.required_exts is not None and item.suffix.lower() not in self.required_exts:
                        continue
                    
                    full_path = str(item.resolve())
                    file_size_bytes = item.stat().st_size
                    mime_type = mimetypes.guess_type(item.name)[0] or "application/octet-stream"
                    
                    file_info = {
                        "type": mime_type,
                        "size_bytes": file_size_bytes
                    }
                    
                    if hasattr(self, 'file_token_counts') and full_path in self.file_token_counts:
                        file_info["token_count"] = self.file_token_counts[full_path]
                        
                    result[item.name] = file_info
                    
            return result
        
        return build_tree(Path(base_path))