fix bug: reader cannot set metadata to DocNode (#540)

Author: wzh1994

lazyllm/tools/rag/dataReader.py

@@ -198,10 +198,11 @@ def _exclude_metadata(self, documents: List[DocNode]) -> List[DocNode]:
     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[AbstractFileSystem] = None,
                   metadata: Optional[Dict] = None) -> List[DocNode]:
-        metadata: dict = metadata or {}
+        # metadata priority: user > reader > metadata_genf
+        user_metadata: Dict = metadata or {}
+        metadata_generated: Dict = metadata_genf(str(input_file)) if metadata_genf is not None else {}
         documents: List[DocNode] = []
 
-        if metadata_genf is not None: metadata.update(metadata_genf(str(input_file)))
         file_reader_patterns = list(file_extractor.keys())
 
         for pattern in file_reader_patterns:
@@ -213,8 +214,11 @@ def load_file(input_file: Path, metadata_genf: Callable[[str], Dict], file_extra
                 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: doc._global_metadata = metadata
-
+                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}"

tests/basic_tests/test_rag_reader.py

@@ -3,6 +3,7 @@
 import pytest
 from lazyllm.tools.rag.readers import ReaderBase
 from lazyllm.tools.rag import SimpleDirectoryReader, DocNode, Document
+from lazyllm.tools.rag.dataReader import RAG_DOC_CREATION_DATE
 
 class YmlReader(ReaderBase):
     def _load_data(self, file, fs=None):
@@ -19,6 +20,13 @@ def processYml(file):
         node._content = "Call the function processYml."
         return [node]
 
+def processYmlWithMetadata(file):
+    with open(file, 'r') as f:
+        data = f.read()
+        node = DocNode(text=data, metadata=dict(m='m'), global_metadata={RAG_DOC_CREATION_DATE: '00-00'})
+        node._content = 'Call the function processYml.'
+        return [node]
+
 class TestRagReader(object):
     def setup_method(self):
         self.doc1 = Document(dataset_path="ci_data/rag_reader_full", manager=False)
@@ -62,6 +70,14 @@ def test_register_global_reader(self):
         docs = self.doc1._impl._reader.load_data(input_files=files)
         assert docs[0].text == "Call the function processYml."
 
+    def test_register_reader_metadata(self):
+        self.doc1.add_reader('**/*.yml', processYmlWithMetadata)
+        files = [os.path.join(self.datasets, 'reader_test.yml')]
+        docs = self.doc1._impl._reader.load_data(input_files=files)
+        assert docs[0].text == 'Call the function processYml.'
+        assert docs[0].metadata.get('m') == 'm'
+        assert docs[0].global_metadata.get(RAG_DOC_CREATION_DATE) == '00-00'
+
     def test_register_local_and_global_reader(self):
         files = [os.path.join(self.datasets, "reader_test.yml")]