|
|
import asyncio |
|
|
from pathlib import Path |
|
|
from typing import Union, List, Tuple, Optional, Callable, Dict |
|
|
|
|
|
from llama_index.core import SimpleDirectoryReader |
|
|
|
|
|
from evoagentx.rag.schema import Document |
|
|
from evoagentx.core.logging import logger |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class LLamaIndexReader: |
|
|
"""A universal file reader based on LlamaIndex's SimpleDirectoryReader. |
|
|
|
|
|
This class provides a flexible interface for loading documents from files or directories, |
|
|
supporting various formats (e.g., PDF, Word, Markdown) with customizable filtering and metadata. |
|
|
|
|
|
Attributes: |
|
|
recursive (bool): Whether to recursively load files from directories. |
|
|
exclude_hidden (bool): Whether to exclude hidden files (starting with '.'). |
|
|
num_workers (Optional[int]): Number of worker threads for parallel loading. |
|
|
num_files_limits (Optional[int]): Maximum number of files to load. |
|
|
custom_metadata_function (Optional[Callable]): Custom function to extract metadata. |
|
|
extern_file_extractor (Optional[Dict]): Custom file extractors for specific file types. |
|
|
errors (str): Error handling strategy for file reading (e.g., 'ignore', 'strict'). |
|
|
encoding (str): File encoding (default: 'utf-8'). |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
recursive: bool = False, |
|
|
exclude_hidden: bool = True, |
|
|
num_workers: Optional[int] = None, |
|
|
num_files_limits: Optional[int] = None, |
|
|
custom_metadata_function: Optional[Callable] = None, |
|
|
extern_file_extractor: Optional[Dict] = None, |
|
|
errors: str = "ignore", |
|
|
encoding: str = "utf-8", |
|
|
): |
|
|
self.recursive = recursive |
|
|
self.exclude_hidden = exclude_hidden |
|
|
self.num_workers = num_workers |
|
|
self.num_files_limits = num_files_limits |
|
|
self.custom_metadata_function = custom_metadata_function |
|
|
self.extern_file_extractor = extern_file_extractor |
|
|
self.errors = errors |
|
|
self.encoding = encoding |
|
|
|
|
|
def _validate_path(self, path: Union[str, Path]) -> Path: |
|
|
"""Validate and convert a path to a Path object. |
|
|
|
|
|
Args: |
|
|
path: A string or Path object representing a file or directory. |
|
|
|
|
|
Returns: |
|
|
Path: A validated Path object. |
|
|
|
|
|
Raises: |
|
|
FileNotFoundError: If the path does not exist. |
|
|
ValueError: If the path is invalid. |
|
|
""" |
|
|
path = Path(path) |
|
|
if not path.exists(): |
|
|
logger.error(f"Path does not exist: {path}") |
|
|
raise FileNotFoundError(f"Path does not exist: {path}") |
|
|
return path |
|
|
|
|
|
def _check_input( |
|
|
self, input_data: Union[str, List, Tuple], is_file: bool = True |
|
|
) -> Union[List[Path], Path]: |
|
|
"""Check input to a list of Path objects or a single Path for directories. |
|
|
|
|
|
Args: |
|
|
input_data: A string, list, or tuple of file/directory paths. |
|
|
is_file: Whether to treat input as file paths (True) or directory (False). |
|
|
|
|
|
Returns: |
|
|
Union[List[Path], Path]: Valied file paths or directory path. |
|
|
|
|
|
Raises: |
|
|
ValueError: If input type is invalid. |
|
|
""" |
|
|
if isinstance(input_data, str): |
|
|
return self._validate_path(input_data) |
|
|
elif isinstance(input_data, (list, tuple)): |
|
|
if is_file: |
|
|
return [self._validate_path(p) for p in input_data] |
|
|
else: |
|
|
return self._validate_path(input_data[0]) |
|
|
else: |
|
|
logger.error(f"Invalid input type: {type(input_data)}") |
|
|
raise ValueError(f"Invalid input type: {type(input_data)}") |
|
|
|
|
|
def load( |
|
|
self, |
|
|
file_paths: Union[str, List, Tuple], |
|
|
exclude_files: Optional[Union[str, List, Tuple]] = None, |
|
|
filter_file_by_suffix: Optional[Union[str, List, Tuple]] = None, |
|
|
merge_by_file: bool = False, |
|
|
show_progress: bool = False, |
|
|
use_async: bool = False, |
|
|
) -> List[Document]: |
|
|
"""Load documents from files or directories. |
|
|
|
|
|
Args: |
|
|
file_paths: A string, list, or tuple of file paths or a directory path. |
|
|
exclude_files: Files to exclude from loading. |
|
|
filter_file_by_suffix: File extensions to include (e.g., ['.pdf', '.docx']). |
|
|
|
|
|
Returns: |
|
|
List[Document]: List of loaded documents. |
|
|
|
|
|
Raises: |
|
|
FileNotFoundError: If input paths are invalid. |
|
|
RuntimeError: If document loading fails. |
|
|
""" |
|
|
try: |
|
|
input_files = None |
|
|
input_dir = None |
|
|
if isinstance(file_paths, (list, tuple)): |
|
|
input_files = self._check_input(file_paths, is_file=True) |
|
|
else: |
|
|
path = self._check_input(file_paths, is_file=False) |
|
|
if path.is_dir(): |
|
|
input_dir = path |
|
|
else: |
|
|
input_files = [path] |
|
|
|
|
|
exclude_files = ( |
|
|
self._check_input(exclude_files, is_file=True) |
|
|
if exclude_files |
|
|
else None |
|
|
) |
|
|
filter_file_by_suffix = ( |
|
|
list(filter_file_by_suffix) |
|
|
if isinstance(filter_file_by_suffix, (list, tuple)) |
|
|
else [filter_file_by_suffix] |
|
|
if isinstance(filter_file_by_suffix, str) |
|
|
else None |
|
|
) |
|
|
|
|
|
reader = SimpleDirectoryReader( |
|
|
input_dir=input_dir, |
|
|
input_files=input_files, |
|
|
exclude=exclude_files, |
|
|
exclude_hidden=self.exclude_hidden, |
|
|
recursive=self.recursive, |
|
|
required_exts=filter_file_by_suffix, |
|
|
num_files_limit=self.num_files_limits, |
|
|
file_metadata=self.custom_metadata_function, |
|
|
file_extractor=self.extern_file_extractor, |
|
|
encoding=self.encoding, |
|
|
errors=self.errors, |
|
|
) |
|
|
|
|
|
llama_docs = asyncio.run(reader.aload_data(show_progress=show_progress, num_workers=self.num_workers)) if use_async \ |
|
|
else reader.load_data(show_progress=show_progress) |
|
|
|
|
|
if merge_by_file: |
|
|
file_to_docs = {} |
|
|
for doc in llama_docs: |
|
|
file_path = doc.metadata.get("file_path", "") |
|
|
if file_path not in file_to_docs: |
|
|
file_to_docs[file_path] = [] |
|
|
file_to_docs[file_path].append(doc) |
|
|
|
|
|
documents = [] |
|
|
for file_path, docs in file_to_docs.items(): |
|
|
combined_text = "\n".join(doc.text for doc in docs) |
|
|
combined = docs[0].copy() |
|
|
combined.text_resource.text = combined_text |
|
|
combined.metadata["page_count"] = len(docs) |
|
|
documents.append(Document.from_llama_document(combined)) |
|
|
else: |
|
|
documents = [Document.from_llama_document(doc) for doc in llama_docs] |
|
|
logger.info(f"Loaded {len(documents)} documents") |
|
|
return documents |
|
|
|
|
|
except Exception as e: |
|
|
logger.error(f"Failed to load documents: {str(e)}") |
|
|
raise RuntimeError(f"Failed to load documents: {str(e)}") |
