Add kwargs

Author: ali-sedaghatbaf

src/neo4j_graphrag/experimental/components/data_loader.py

@@ -19,7 +19,7 @@
 import io
 from abc import abstractmethod
 from pathlib import Path
-from typing import Dict, Optional, Union, cast
+from typing import Any, Dict, Optional, Union, cast
 
 import fsspec
 import pypdf
@@ -49,21 +49,12 @@ 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,
+        **kwargs: Any,
     ) -> LoadedDocument: ...
 
 
@@ -92,18 +83,16 @@ 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,
+        **kwargs: Any,
     ) -> LoadedDocument:
+        fs = kwargs.get("fs")
         if not isinstance(filepath, str):
             filepath = str(filepath)
         if isinstance(fs, str):
             fs = fsspec.filesystem(fs)
         elif fs is None:
             fs = LocalFileSystem()
-        text = self._apply_max_chars(
-            self.load_file(filepath, fs), max_chars=max_chars
-        )
+        text = self.load_file(filepath, fs)
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(
@@ -133,18 +122,16 @@ 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,
+        **kwargs: Any,
     ) -> LoadedDocument:
+        fs = kwargs.get("fs")
         if not isinstance(filepath, str):
             filepath = str(filepath)
         if isinstance(fs, str):
             fs = fsspec.filesystem(fs)
         elif fs is None:
             fs = LocalFileSystem()
-        text = self._apply_max_chars(
-            MarkdownLoader.load_file(filepath, fs), max_chars=max_chars
-        )
+        text = MarkdownLoader.load_file(filepath, fs)
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(

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

@@ -26,7 +26,6 @@
 )
 import warnings
 
-from fsspec import AbstractFileSystem
 from pydantic import ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 
@@ -84,24 +83,22 @@ 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,
+        **kwargs: Any,
     ) -> LoadedDocument:
+        fs = kwargs.get("fs")
         path_str = str(filepath)
         suffix = Path(path_str).suffix.lower()
         if suffix == ".pdf":
             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,
-                max_chars=max_chars,
             )
         raise UnsupportedDocumentFormatError(
             f"Unsupported document format: {suffix!r}. "
@@ -439,10 +436,6 @@ 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,14 +87,6 @@ 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``)."""
@@ -106,14 +98,6 @@ 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,26 +370,6 @@ 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"})