Add max_chars

Author: ali-sedaghatbaf

src/neo4j_graphrag/experimental/components/data_loader.py

@@ -49,11 +49,21 @@ def get_document_metadata(
     ) -> Dict[str, str] | None:
         return metadata
 
+    @staticmethod
+    def _apply_max_chars(text: str, max_chars: Optional[int] = None) -> str:
+        if max_chars is None:
+            return text
+        if max_chars < 0:
+            raise ValueError("max_chars must be >= 0")
+        return text[:max_chars]
+
     @abstractmethod
     async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument: ...
 
 
@@ -83,14 +93,17 @@ async def run(
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
         fs: Optional[Union[AbstractFileSystem, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
         if not isinstance(filepath, str):
             filepath = str(filepath)
         if isinstance(fs, str):
             fs = fsspec.filesystem(fs)
         elif fs is None:
             fs = LocalFileSystem()
-        text = self.load_file(filepath, fs)
+        text = self._apply_max_chars(
+            self.load_file(filepath, fs), max_chars=max_chars
+        )
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(
@@ -121,14 +134,17 @@ async def run(
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
         fs: Optional[Union[AbstractFileSystem, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
         if not isinstance(filepath, str):
             filepath = str(filepath)
         if isinstance(fs, str):
             fs = fsspec.filesystem(fs)
         elif fs is None:
             fs = LocalFileSystem()
-        text = MarkdownLoader.load_file(filepath, fs)
+        text = self._apply_max_chars(
+            MarkdownLoader.load_file(filepath, fs), max_chars=max_chars
+        )
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(

src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py

@@ -85,14 +85,23 @@ async def run(
         filepath: Union[str, Path],
         metadata: Optional[dict[str, str]] = None,
         fs: Optional[Union[AbstractFileSystem, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
         path_str = str(filepath)
         suffix = Path(path_str).suffix.lower()
         if suffix == ".pdf":
-            return await PdfLoader().run(filepath=path_str, metadata=metadata, fs=fs)
+            return await PdfLoader().run(
+                filepath=path_str,
+                metadata=metadata,
+                fs=fs,
+                max_chars=max_chars,
+            )
         if suffix in (".md", ".markdown"):
             return await MarkdownLoader().run(
-                filepath=path_str, metadata=metadata, fs=fs
+                filepath=path_str,
+                metadata=metadata,
+                fs=fs,
+                max_chars=max_chars,
             )
         raise UnsupportedDocumentFormatError(
             f"Unsupported document format: {suffix!r}. "
@@ -430,6 +439,10 @@ def get_run_params(self, user_input: dict[str, Any]) -> dict[str, Any]:
                 )
             run_params["file_loader"]["filepath"] = file_path
             run_params["file_loader"]["metadata"] = user_input.get("document_metadata")
+            max_chars = user_input.get("max_chars")
+            # Backward-compatible: only forward new arg for the default loader.
+            if max_chars is not None and self.file_loader is None:
+                run_params["file_loader"]["max_chars"] = max_chars
         else:
             if not text:
                 raise PipelineDefinitionError(

tests/unit/experimental/components/test_data_loader.py

@@ -87,6 +87,14 @@ async def test_markdown_loader_run() -> None:
     assert "# Hello" in doc.text
 
 
+@pytest.mark.asyncio
+async def test_markdown_loader_run_max_chars() -> None:
+    md_path = str(BASE_DIR / "sample_data/hello.md")
+    loader = MarkdownLoader()
+    doc = await loader.run(filepath=md_path, max_chars=7)
+    assert doc.text == "# Hello"
+
+
 @pytest.mark.asyncio
 async def test_pdf_loader_run() -> None:
     """``PdfLoader.run`` wraps ``load_file`` with :class:`DocumentInfo` (default ``fs``)."""
@@ -98,6 +106,14 @@ async def test_pdf_loader_run() -> None:
     assert doc.text == "Lorem ipsum dolor sit amet."
 
 
+@pytest.mark.asyncio
+async def test_pdf_loader_run_max_chars() -> None:
+    pdf_path = str(BASE_DIR / "sample_data/lorem_ipsum.pdf")
+    loader = PdfLoader()
+    doc = await loader.run(filepath=pdf_path, max_chars=5)
+    assert doc.text == "Lorem"
+
+
 @pytest.mark.asyncio
 async def test_pdf_loader_run_fs_string_resolves_with_fsspec(
     dummy_pdf_path: str,

tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py

@@ -370,6 +370,26 @@ def test_simple_kg_pipeline_config_run_params_from_file_file_path() -> None:
     }
 
 
+def test_simple_kg_pipeline_config_run_params_from_file_file_path_with_max_chars() -> (
+    None
+):
+    config = SimpleKGPipelineConfig(from_file=True)
+    assert config.get_run_params({"file_path": "my_file", "max_chars": 42}) == {
+        "file_loader": {"filepath": "my_file", "metadata": None, "max_chars": 42}
+    }
+
+
+def test_simple_kg_pipeline_config_run_params_custom_file_loader_ignores_max_chars() -> (
+    None
+):
+    config = SimpleKGPipelineConfig(
+        from_file=True, file_loader=ComponentType(PdfLoader())
+    )
+    assert config.get_run_params({"file_path": "my_file", "max_chars": 42}) == {
+        "file_loader": {"filepath": "my_file", "metadata": None}
+    }
+
+
 def test_simple_kg_pipeline_config_run_params_from_text_text() -> None:
     config = SimpleKGPipelineConfig(from_file=False)
     run_params = config.get_run_params({"text": "my text"})