Add max_chars

Author: ali-sedaghatbaf

examples/customize/build_graph/components/loaders/custom_loader.py

@@ -12,8 +12,10 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
         # Implement logic here to read and transform the input file.
+        _ = max_chars
         return LoadedDocument(
             text="<extracted text>",
             document_info=DocumentInfo(

src/neo4j_graphrag/experimental/components/data_loader.py

@@ -41,11 +41,20 @@ 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,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument: ...
 
 
@@ -70,10 +79,11 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
         if not isinstance(filepath, str):
             filepath = str(filepath)
-        text = self.load_file(filepath)
+        text = self._apply_max_chars(self.load_file(filepath), max_chars=max_chars)
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(
@@ -100,10 +110,13 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
         if not isinstance(filepath, str):
             filepath = str(filepath)
-        text = MarkdownLoader.load_file(filepath)
+        text = self._apply_max_chars(
+            MarkdownLoader.load_file(filepath), max_chars=max_chars
+        )
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(

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

@@ -83,13 +83,18 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[dict[str, 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)
+            return await PdfLoader().run(
+                filepath=path_str, metadata=metadata, max_chars=max_chars
+            )
         if suffix in (".md", ".markdown"):
-            return await MarkdownLoader().run(filepath=path_str, metadata=metadata)
+            return await MarkdownLoader().run(
+                filepath=path_str, metadata=metadata, max_chars=max_chars
+            )
         raise UnsupportedDocumentFormatError(
             f"Unsupported document format: {suffix!r}. "
             f"Supported: .pdf, .md, .markdown"
@@ -426,6 +431,9 @@ 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")
+            if max_chars is not None:
+                run_params["file_loader"]["max_chars"] = max_chars
         else:
             if not text:
                 raise PipelineDefinitionError(

tests/unit/experimental/components/test_data_loader.py

@@ -79,6 +79,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 = 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`."""
@@ -90,6 +98,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 = 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_run_passes_metadata_to_document_info(dummy_pdf_path: str) -> None:
     loader = PdfLoader()
@@ -105,8 +121,11 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[dict[str, str]] = None,
+        max_chars: Optional[int] = None,
     ) -> LoadedDocument:
-        return await super().run(filepath=filepath, metadata=metadata)
+        return await super().run(
+            filepath=filepath, metadata=metadata, max_chars=max_chars
+        )
 
     def get_document_metadata(
         self, text: str, metadata: dict[str, str] | None = None

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

@@ -370,6 +370,15 @@ 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": 32}) == {
+        "file_loader": {"filepath": "my_file", "metadata": None, "max_chars": 32}
+    }
+
+
 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"})