support more default node types (#789)

Author: wzh1994

docs/en/API Reference/tools.md

@@ -229,6 +229,7 @@
 
 ::: lazyllm.tools.rag.doc_processor.DocumentProcessor
     members: register_algorithm, drop_algorithm
+
 ::: lazyllm.tools.rag.doc_node.QADocNode
     members: get_text
     exclude-members:

lazyllm/docs/tools.py

@@ -2620,7 +2620,7 @@ def my_reranker(node: DocNode, **kwargs):
 Args:
     name (str): 算法名称，作为唯一标识符。
     store (StoreBase): 存储实例，用于管理文档数据。
-    reader (ReaderBase): 读取器实例，用于解析文档内容。
+    reader (DirectoryReader): 读取器实例，用于解析文档内容。
     node_groups (Dict[str, Dict]): 节点组配置信息。
     display_name (Optional[str]): 算法的显示名称，默认为None。
     description (Optional[str]): 算法的描述信息，默认为None。
@@ -2638,7 +2638,7 @@ def my_reranker(node: DocNode, **kwargs):
 Args:
     name (str): Algorithm name as unique identifier.
     store (StoreBase): Storage instance for managing document data.
-    reader (ReaderBase): Reader instance for parsing document content.
+    reader (DirectoryReader): Reader instance for parsing document content.
     node_groups (Dict[str, Dict]): Node group configuration information.
     display_name (Optional[str]): Display name for the algorithm, defaults to None.
     description (Optional[str]): Description of the algorithm, defaults to None.
@@ -2652,11 +2652,11 @@ def my_reranker(node: DocNode, **kwargs):
 
 add_example('rag.doc_processor.DocumentProcessor.register_algorithm', """
 ```python
-from lazyllm.rag import DocumentProcessor, FileStore, PDFReader
+from lazyllm.rag import DocumentProcessor, FileStore, DirectoryReader
 
 # Create storage and reader instances
 store = FileStore(path="./data")
-reader = PDFReader()
+reader = DirectoryReader()
 
 # Define node group configuration
 node_groups = {
@@ -3563,10 +3563,10 @@ def _lazy_load_data(self, file_paths: list, **kwargs) -> Iterable[DocNode]:
 Args:
     input_files (Optional[List[str]]): A list of file paths to read. If None, uses the files specified during initialization.
     metadatas (Optional[Dict]): Additional metadata to associate with the loaded documents.
-    split_image_nodes (bool): Whether to separate image nodes from text nodes. If True, returns a tuple of (text_nodes, image_nodes). If False, returns all nodes together.
+    split_nodes_by_type (bool): Whether to separate image and other nodes from text nodes. If True, returns a tuple of (text_nodes, image_nodes). If False, returns all nodes together.
 
 **Returns:**\n
-- Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: If split_image_nodes is False, returns a list of all document nodes. If True, returns a tuple containing text nodes and image nodes separately.
+- Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: If split_nodes_by_type is False, returns a list of all document nodes. If True, returns a tuple containing text nodes and image nodes separately.
 ''')
 
 add_chinese_doc('rag.data_loaders.DirectoryReader.load_data', '''\
@@ -3577,10 +3577,10 @@ def _lazy_load_data(self, file_paths: list, **kwargs) -> Iterable[DocNode]:
 Args:
     input_files (Optional[List[str]]): 要读取的文件路径列表。如果为None，使用初始化时指定的文件。
     metadatas (Optional[Dict]): 与加载文档关联的额外元数据。
-    split_image_nodes (bool): 是否将图像节点与文本节点分离。如果为True，返回(text_nodes, image_nodes)的元组。如果为False，一起返回所有节点。
+    split_nodes_by_type (bool): 是否将图像等其他节点与文本节点分离。如果为True，返回(text_nodes, image_nodes)的元组。如果为False，一起返回所有节点。
 
 **Returns:**\n
-- Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: 如果split_image_nodes为False，返回所有文档节点的列表。如果为True，返回包含文本节点和图像节点的元组。
+- Union[List[DocNode], Tuple[List[DocNode], List[ImageDocNode]]]: 如果split_nodes_by_type为False，返回所有文档节点的列表。如果为True，返回包含文本节点和图像节点的元组。
 ''')
 
 # ---------------------------------------------------------------------------- #

lazyllm/tools/rag/__init__.py

@@ -13,7 +13,11 @@
 from .index_base import IndexBase
 from .store import LazyLLMStoreBase
 
+
+add_post_action_for_default_reader = SimpleDirectoryReader.add_post_action_for_default_reader
+
 __all__ = [
+    'add_post_action_for_default_reader',
     'Document',
     'Reranker',
     'Retriever',

lazyllm/tools/rag/dataReader.py

@@ -6,19 +6,22 @@
 import mimetypes
 import multiprocessing
 import fnmatch
+import traceback
 from tqdm import tqdm
 from datetime import datetime
 from functools import reduce
 from itertools import repeat
-from typing import Dict, Optional, List, Callable, Type
+from typing import Dict, Optional, List, Callable, Type, Union
 from pathlib import Path, PurePosixPath, PurePath
 from lazyllm.thirdparty import fsspec
 from lazyllm import ModuleBase, LOG, config
 from lazyllm.components.formatter.formatterbase import _lazyllm_get_file_list
+from lazyllm.tools.rag.readers.readerBase import TxtReader, DefaultReader
 from .doc_node import DocNode
 from .readers import (ReaderBase, PDFReader, DocxReader, HWPReader, PPTXReader, ImageReader, IPYNBReader,
                       EpubReader, MarkdownReader, MboxReader, PandasCSVReader, PandasExcelReader, VideoAudioReader,
                       get_default_fs, is_default_fs)
+from .transform import NodeTransform, FuncNodeTransform
 from .global_metadata import (RAG_DOC_PATH, RAG_DOC_FILE_NAME, RAG_DOC_FILE_TYPE, RAG_DOC_FILE_SIZE,
                               RAG_DOC_CREATION_DATE, RAG_DOC_LAST_MODIFIED_DATE, RAG_DOC_LAST_ACCESSED_DATE)
 
@@ -79,6 +82,8 @@ class SimpleDirectoryReader(ModuleBase):
         '*.xlsx': PandasExcelReader,
         '*.mp3': VideoAudioReader,
         '*.mp4': VideoAudioReader,
+        '*.txt': TxtReader,
+        '*.xml': TxtReader,
     }
 
     def __init__(self, input_dir: Optional[str] = None, input_files: Optional[List] = None,
@@ -89,39 +94,31 @@ def __init__(self, input_dir: Optional[str] = None, input_files: Optional[List]
                  return_trace: bool = False, metadatas: Optional[Dict] = None) -> None:
         super().__init__(return_trace=return_trace)
 
-        if (not input_dir and not input_files) or (input_dir and input_files):
-            raise ValueError('Must provide either `input_dir` or `input_files`.')
-
         self._fs = fs or get_default_fs()
         self._encoding = encoding
-
         self._exclude = exclude
         self._recursive = recursive
         self._exclude_hidden = exclude_hidden
         self._required_exts = required_exts
         self._num_files_limit = num_files_limit
         self._Path = Path if is_default_fs(self._fs) else PurePosixPath
         self._metadatas = metadatas
+        self._input_files = self._get_input_files(input_dir, input_files)
+        self._file_extractor = {**self.default_file_readers, **(file_extractor or {})}
+        self._metadata_genf = metadata_genf or _DefaultFileMetadataFunc(self._fs)
+        if filename_as_id: LOG.warning('Argument `filename_as_id` for DataReader is no longer used')
 
+    def _get_input_files(self, input_dir, input_files):
         if input_files:
-            self._input_files = []
-            for path in input_files:
-                if not self._fs.isfile(path):
-                    path = os.path.join(config['data_path'], path)
-                    if not self._fs.isfile(path):
-                        raise ValueError(f'File {path} does not exist.')
-                input_file = self._Path(path)
-                self._input_files.append(input_file)
+            assert not input_dir, 'Cannot provide files and dir at the same time'
+            input_files = [os.path.join(config['data_path'], p) if not self._fs.isfile(p) else p for p in input_files]
+            input_files = [self._Path(p) if p else (_ for _ in ()).throw(ValueError, f'File {p} does not exist.')
+                           for p in input_files]
         elif input_dir:
             if not self._fs.isdir(input_dir):
                 raise ValueError(f'Directory {input_dir} does not exist.')
-            self._input_dir = self._Path(input_dir)
-            self._input_files = self._add_files(self._input_dir)
-
-        self._file_extractor = file_extractor or {}
-
-        self._metadata_genf = metadata_genf or _DefaultFileMetadataFunc(self._fs)
-        if filename_as_id: LOG.warning('Argument `filename_as_id` for DataReader is no longer used')
+            input_files = self._add_files(self._Path(input_dir))
+        return input_files
 
     def _add_files(self, input_dir: Path) -> List[Path]:  # noqa: C901
         all_files = set()
@@ -195,15 +192,10 @@ def _exclude_metadata(self, documents: List[DocNode]) -> List[DocNode]:
         return documents
 
     @staticmethod
-    def load_file(input_file: Path, metadata_genf: Callable[[str], Dict], file_extractor: Dict[str, Callable],
-                  encoding: str = 'utf-8', pathm: PurePath = Path, fs: Optional['fsspec.AbstractFileSystem'] = None,
-                  metadata: Optional[Dict] = None) -> List[DocNode]:
-        # metadata priority: user > reader > metadata_genf
-        user_metadata: Dict = metadata or {}
-        metadata_generated: Dict = metadata_genf(str(input_file)) if metadata_genf else {}
-        documents: List[DocNode] = []
-
+    def find_extractor_by_file(input_file: Path, file_extractor: Dict[str, Callable], pathm: PurePath = Path):
         filename_lower = str(input_file).lower()
+        file_suffix = filename_lower.split('.')[-1]
+        if extractor := file_extractor.get(f'*.{file_suffix}'): return extractor
 
         for pattern, extractor in file_extractor.items():
             pt_lower = str(pathm(pattern)).lower()
@@ -213,72 +205,91 @@ def load_file(input_file: Path, metadata_genf: Callable[[str], Dict], file_extra
             else:
                 base = str(pathm.cwd()).lower()
                 match_pattern = os.path.join(base, pt_lower)
-
             if fnmatch.fnmatch(filename_lower, match_pattern):
-                reader = extractor() if isinstance(extractor, type) else extractor
-                kwargs = {'fs': fs} if fs and not is_default_fs(fs) else {}
-                docs = reader(input_file, **kwargs)
-                if isinstance(docs, DocNode): docs = [docs]
-                for doc in docs:
-                    metadata = metadata_generated.copy()
-                    metadata.update(doc._global_metadata or {})
-                    metadata.update(user_metadata)
-                    doc._global_metadata = metadata
-
-                if config['rag_filename_as_id']:
-                    for i, doc in enumerate(docs):
-                        doc._uid = f'{input_file!s}_index_{i}'
-                documents.extend(docs)
-                break
-        else:
-            if not config['use_fallback_reader']:
-                LOG.warning(f'no pattern found for {input_file}! '
-                            'If you want fallback to default Reader, set `LAZYLLM_USE_FALLBACK_READER=True`.')
-                return documents
-            fs = fs or get_default_fs()
-            with fs.open(input_file, encoding=encoding) as f:
-                try:
-                    data = f.read().decode(encoding)
-                    doc = DocNode(text=data, global_metadata=user_metadata)
-                    documents.append(doc)
-                except Exception:
-                    LOG.error(f'no pattern found for {input_file} and it is not utf-8, skip it!')
-        return documents
+                return extractor
+        return DefaultReader
 
-    def _load_data(self, show_progress: bool = False, num_workers: Optional[int] = None,
-                   fs: Optional['fsspec.AbstractFileSystem'] = None) -> List[DocNode]:
-        documents = []
+    @staticmethod
+    def load_file(input_file: Path, metadata_genf: Callable[[str], Dict], file_extractor: Dict[str, Callable],
+                  encoding: str = 'utf-8', pathm: PurePath = Path, fs: Optional['fsspec.AbstractFileSystem'] = None,
+                  metadata: Optional[Dict] = None) -> List[DocNode]:
+        # metadata priority: user > reader > metadata_genf
+        user_metadata: Dict = metadata or {}
+        metadata_generated: Dict = metadata_genf(str(input_file)) if metadata_genf else {}
+        rd = SimpleDirectoryReader.find_extractor_by_file(input_file, file_extractor, pathm)
+        reader = rd(encoding=encoding) if isinstance(rd, TxtReader) else rd() if isinstance(rd, type) else rd
+        kwargs = {'fs': fs} if fs and not is_default_fs(fs) else {}
+
+        try:
+            docs = reader(input_file, **kwargs)
+        except Exception as e:
+            LOG.error(f'Error loading file {input_file}, skip it!')
+            LOG.error(f'message: {e}\n Traceback: {traceback.format_tb(e.__traceback__)}')
+            return []
+        docs = [docs] if isinstance(docs, DocNode) else [] if docs is None else docs
+
+        for doc in docs:
+            metadata = metadata_generated.copy()
+            metadata.update(doc._global_metadata or {})
+            metadata.update(user_metadata)
+            doc._global_metadata = metadata
 
-        fs = fs or self._fs
-        process_file = self._input_files
-        file_readers = self._file_extractor.copy()
-        for key, func in self.default_file_readers.items():
-            if key not in file_readers: file_readers[key] = func
+        if config['rag_filename_as_id']:
+            for i, doc in enumerate(docs):
+                doc._uid = f'{input_file!s}_index_{i}'
+        return docs
+
+    def _load_data(self, show_progress: bool = False, num_workers: Optional[int] = None,
+                   fs: Optional['fsspec.AbstractFileSystem'] = None, metadatas: Optional[Dict] = None,
+                   input_dir: Optional[str] = None, input_files: Optional[List] = None) -> List[DocNode]:
+        documents, fs, metadatas = [], fs or self._fs, metadatas or self._metadatas
+        process_file = self._get_input_files(input_dir, input_files) if input_dir or input_files else self._input_files
 
         if num_workers and num_workers >= 1:
             if num_workers > multiprocessing.cpu_count():
                 LOG.warning('Specified num_workers exceed number of CPUs in the system. '
                             'Setting `num_workers` down to the maximum CPU count.')
             with multiprocessing.get_context('spawn').Pool(num_workers) as p:
                 results = p.starmap(SimpleDirectoryReader.load_file,
-                                    zip(process_file, repeat(self._metadata_genf), repeat(file_readers),
+                                    zip(process_file, repeat(self._metadata_genf), repeat(self._file_extractor),
                                         repeat(self._encoding), repeat(self._Path),
-                                        repeat(self._fs), self._metadatas or repeat(None)))
+                                        repeat(self._fs), metadatas or repeat(None)))
                 documents = reduce(lambda x, y: x + y, results)
         else:
             if show_progress:
                 process_file = tqdm(self._input_files, desc='Loading files', unit='file')
-            for input_file, metadata in zip(process_file, self._metadatas or repeat(None)):
+            for input_file, metadata in zip(process_file, metadatas or repeat(None)):
                 documents.extend(
                     SimpleDirectoryReader.load_file(
-                        input_file=input_file, metadata_genf=self._metadata_genf, file_extractor=file_readers,
+                        input_file=input_file, metadata_genf=self._metadata_genf, file_extractor=self._file_extractor,
                         encoding=self._encoding, pathm=self._Path, fs=self._fs, metadata=metadata))
 
         return self._exclude_metadata(documents)
 
     def forward(self, *args, **kwargs) -> List[DocNode]:
         return self._load_data(*args, **kwargs)
 
+    @staticmethod
+    def get_default_reader(file_ext: str) -> Callable[[Path, Dict], List[DocNode]]:
+        if not file_ext.startswith('*.'): file_ext = '*.' + file_ext
+        return SimpleDirectoryReader.default_file_readers.get(file_ext)
+
+    @staticmethod
+    def add_post_action_for_default_reader(file_ext: str, f: Callable[[DocNode], Union[DocNode, List[DocNode]]]):
+        if not file_ext.startswith('*.'): file_ext = '*.' + file_ext
+        if file_ext not in SimpleDirectoryReader.default_file_readers:
+            raise KeyError(f'{file_ext} has no default reader, use Document.add_reader instead')
+
+        reader = SimpleDirectoryReader.default_file_readers[file_ext]
+        assert isinstance(reader, type) and issubclass(reader, ReaderBase)
+
+        if isinstance(f, type): f = f()
+        if not isinstance(f, NodeTransform):
+            try: f('test')
+            except Exception: pass
+            else: f = FuncNodeTransform(f, trans_node=False)
+        reader.post_action = staticmethod(f)
+
 
 config.add('rag_filename_as_id', bool, False, 'RAG_FILENAME_AS_ID')
 config.add('use_fallback_reader', bool, True, 'USE_FALLBACK_READER')

lazyllm/tools/rag/data_loaders.py

@@ -1,36 +1,36 @@
-from typing import List, Optional, Dict
+from typing import List, Optional, Dict, Union
 from lazyllm import LOG
+from lazyllm.common.common import once_wrapper
 
 from .doc_node import DocNode, ImageDocNode
 from .store import LAZY_ROOT_NAME, LAZY_IMAGE_GROUP
 from .dataReader import SimpleDirectoryReader
+from collections import defaultdict
+
+type_mapping = {
+    DocNode: LAZY_ROOT_NAME,
+    ImageDocNode: LAZY_IMAGE_GROUP,
+}
 
 class DirectoryReader:
     def __init__(self, input_files: Optional[List[str]], local_readers: Optional[Dict] = None,
                  global_readers: Optional[Dict] = None) -> None:
         self._input_files = input_files
-        self._local_readers = local_readers
-        self._global_readers = global_readers
+        self._local_readers, self._global_readers = local_readers, global_readers
+
+    @once_wrapper
+    def _lazy_init(self):
+        self._reader = SimpleDirectoryReader(file_extractor={**self._global_readers, **self._local_readers})
 
     def load_data(self, input_files: Optional[List[str]] = None, metadatas: Optional[Dict] = None,
-                  *, split_image_nodes: bool = False) -> List[DocNode]:
+                  *, split_nodes_by_type: bool = False) -> List[DocNode]:
+        self._lazy_init()
         input_files = input_files or self._input_files
-        file_readers = self._local_readers.copy()
-        for key, func in self._global_readers.items():
-            if key not in file_readers: file_readers[key] = func
-        LOG.info(f'DirectoryReader loads data, input files: {input_files}')
-        reader = SimpleDirectoryReader(input_files=input_files, file_extractor=file_readers, metadatas=metadatas)
-        nodes: List[DocNode] = []
-        image_nodes: List[ImageDocNode] = []
-        for doc in reader():
-            doc._group = LAZY_IMAGE_GROUP if isinstance(doc, ImageDocNode) else LAZY_ROOT_NAME
-            if not split_image_nodes or not isinstance(doc, ImageDocNode):
-                nodes.append(doc)
-            else:
-                image_nodes.append(doc)
-        if not nodes and not image_nodes:
-            LOG.warning(
-                f'No nodes load from path {input_files}, please check your data path.'
-            )
+        nodes: Union[List[DocNode], Dict[str, List[DocNode]]] = defaultdict(list) if split_nodes_by_type else []
+        for doc in self._reader(input_files=input_files, metadatas=metadatas):
+            doc._group = type_mapping.get(type(doc), LAZY_ROOT_NAME)
+            nodes[doc._group].append(doc) if split_nodes_by_type else nodes.append(doc)
+        if not nodes:
+            LOG.warning(f'No nodes load from path {input_files}, please check your data path.')
         LOG.info('DirectoryReader loads data done!')
-        return (nodes, image_nodes) if split_image_nodes else nodes
+        return nodes