Revert "Remove fs"

This reverts commit baa11ee.

Author: ali-sedaghatbaf

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

@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
+from fsspec import AbstractFileSystem
+
 from neo4j_graphrag.experimental.components.data_loader import DataLoader
 from neo4j_graphrag.experimental.components.types import DocumentInfo, LoadedDocument
 
@@ -12,8 +14,10 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
     ) -> LoadedDocument:
-        # Implement logic here to read and transform the input file.
+        # Implement logic here; use ``fs`` when reading from non-local storage.
+        _ = fs
         return LoadedDocument(
             text="<extracted text>",
             document_info=DocumentInfo(

src/neo4j_graphrag/experimental/components/data_loader.py

@@ -16,11 +16,15 @@
 
 from __future__ import annotations
 
+import io
 from abc import abstractmethod
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Union, cast
 
+import fsspec
 import pypdf
+from fsspec import AbstractFileSystem
+from fsspec.implementations.local import LocalFileSystem
 
 from neo4j_graphrag.exceptions import MarkdownLoadError, PdfLoaderError
 from neo4j_graphrag.experimental.components.types import (
@@ -31,6 +35,10 @@
 from neo4j_graphrag.experimental.pipeline.component import Component
 
 
+def is_default_fs(fs: fsspec.AbstractFileSystem) -> bool:
+    return isinstance(fs, LocalFileSystem) and not fs.auto_mkdir
+
+
 class DataLoader(Component):
     """
     Interface for loading data of various input types.
@@ -46,18 +54,23 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
     ) -> LoadedDocument: ...
 
 
 class PdfLoader(DataLoader):
     """Loads text from PDF files using pypdf."""
 
     @staticmethod
-    def load_file(file: str) -> str:
+    def load_file(
+        file: str,
+        fs: AbstractFileSystem,
+    ) -> str:
         """Parse a PDF file and return extracted text."""
         try:
-            with open(file, "rb") as fp:
-                pdf = pypdf.PdfReader(fp)
+            with fs.open(file, "rb") as fp:
+                stream = fp if is_default_fs(fs) else io.BytesIO(fp.read())
+                pdf = pypdf.PdfReader(stream)
                 num_pages = len(pdf.pages)
                 text_parts = (
                     pdf.pages[page].extract_text() for page in range(num_pages)
@@ -70,10 +83,15 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
     ) -> LoadedDocument:
         if not isinstance(filepath, str):
             filepath = str(filepath)
-        text = self.load_file(filepath)
+        if isinstance(fs, str):
+            fs = fsspec.filesystem(fs)
+        elif fs is None:
+            fs = LocalFileSystem()
+        text = self.load_file(filepath, fs)
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(
@@ -88,22 +106,30 @@ class MarkdownLoader(DataLoader):
     """Loads UTF-8 Markdown (``.md`` / ``.markdown``) files as plain text."""
 
     @staticmethod
-    def load_file(file: str) -> str:
+    def load_file(
+        file: str,
+        fs: AbstractFileSystem,
+    ) -> str:
         try:
-            with open(file, "rb") as fp:
+            with fs.open(file, "rb") as fp:
                 raw = fp.read()
-            return raw.decode("utf-8")
+            return cast(str, raw.decode("utf-8"))
         except Exception as e:
             raise MarkdownLoadError(e)
 
     async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
     ) -> LoadedDocument:
         if not isinstance(filepath, str):
             filepath = str(filepath)
-        text = MarkdownLoader.load_file(filepath)
+        if isinstance(fs, str):
+            fs = fsspec.filesystem(fs)
+        elif fs is None:
+            fs = LocalFileSystem()
+        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,6 +26,7 @@
 )
 import warnings
 
+from fsspec import AbstractFileSystem
 from pydantic import ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 
@@ -83,13 +84,16 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = 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, fs=fs)
         if suffix in (".md", ".markdown"):
-            return await MarkdownLoader().run(filepath=path_str, metadata=metadata)
+            return await MarkdownLoader().run(
+                filepath=path_str, metadata=metadata, fs=fs
+            )
         raise UnsupportedDocumentFormatError(
             f"Unsupported document format: {suffix!r}. "
             f"Supported: .pdf, .md, .markdown"

tests/unit/experimental/components/test_data_loader.py

@@ -20,6 +20,8 @@
 from unittest.mock import patch
 
 import pytest
+from fsspec import AbstractFileSystem
+from fsspec.implementations.local import LocalFileSystem
 from neo4j_graphrag.exceptions import MarkdownLoadError, PdfLoaderError
 from neo4j_graphrag.experimental.components.data_loader import (
     MarkdownLoader,
@@ -47,25 +49,31 @@ def dummy_md_path() -> str:
 
 def test_pdf_loading(pdf_loader: PdfLoader, dummy_pdf_path: str) -> None:
     expected_content = "Lorem ipsum dolor sit amet."
-    actual_content = pdf_loader.load_file(dummy_pdf_path)
+    actual_content = pdf_loader.load_file(dummy_pdf_path, fs=LocalFileSystem())
     assert actual_content == expected_content
 
 
 def test_pdf_processing_error(pdf_loader: PdfLoader, dummy_pdf_path: str) -> None:
-    with patch("builtins.open", side_effect=Exception("Failed to open")):
+    with patch(
+        "fsspec.implementations.local.LocalFileSystem.open",
+        side_effect=Exception("Failed to open"),
+    ):
         with pytest.raises(PdfLoaderError):
-            pdf_loader.load_file(dummy_pdf_path)
+            pdf_loader.load_file(dummy_pdf_path, fs=LocalFileSystem())
 
 
 def test_markdown_processing_error(dummy_md_path: str) -> None:
-    with patch("builtins.open", side_effect=Exception("Failed to open")):
+    with patch(
+        "fsspec.implementations.local.LocalFileSystem.open",
+        side_effect=Exception("Failed to open"),
+    ):
         with pytest.raises(MarkdownLoadError):
-            MarkdownLoader.load_file(dummy_md_path)
+            MarkdownLoader.load_file(dummy_md_path, fs=LocalFileSystem())
 
 
 def test_markdown_loading() -> None:
     md_path = str(BASE_DIR / "sample_data/hello.md")
-    text = MarkdownLoader.load_file(md_path)
+    text = MarkdownLoader.load_file(md_path, fs=LocalFileSystem())
     assert "# Hello" in text
     assert "Markdown **content**" in text
 
@@ -81,7 +89,7 @@ async def test_markdown_loader_run() -> None:
 
 @pytest.mark.asyncio
 async def test_pdf_loader_run() -> None:
-    """``PdfLoader.run`` wraps ``load_file`` with :class:`DocumentInfo`."""
+    """``PdfLoader.run`` wraps ``load_file`` with :class:`DocumentInfo` (default ``fs``)."""
     pdf_path = BASE_DIR / "sample_data/lorem_ipsum.pdf"
     loader = PdfLoader()
     doc = await loader.run(filepath=pdf_path)
@@ -90,6 +98,25 @@ async def test_pdf_loader_run() -> None:
     assert doc.text == "Lorem ipsum dolor sit amet."
 
 
+@pytest.mark.asyncio
+async def test_pdf_loader_run_fs_string_resolves_with_fsspec(
+    dummy_pdf_path: str,
+) -> None:
+    """``fs`` may be a protocol name passed to ``fsspec.filesystem`` (e.g. ``\"file\"``)."""
+    loader = PdfLoader()
+    doc = await loader.run(filepath=dummy_pdf_path, fs="file")
+    assert "Lorem ipsum" in doc.text
+
+
+@pytest.mark.asyncio
+async def test_markdown_loader_run_fs_string() -> None:
+    md_path = str(BASE_DIR / "sample_data/hello.md")
+    loader = MarkdownLoader()
+    doc = await loader.run(filepath=md_path, fs="file")
+    assert doc.document_info.document_type == DocumentType.MARKDOWN
+    assert "# Hello" in doc.text
+
+
 @pytest.mark.asyncio
 async def test_run_passes_metadata_to_document_info(dummy_pdf_path: str) -> None:
     loader = PdfLoader()
@@ -105,8 +132,9 @@ async def run(
         self,
         filepath: Union[str, Path],
         metadata: Optional[dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
     ) -> LoadedDocument:
-        return await super().run(filepath=filepath, metadata=metadata)
+        return await super().run(filepath=filepath, metadata=metadata, fs=fs)
 
     def get_document_metadata(
         self, text: str, metadata: dict[str, str] | None = None
@@ -130,6 +158,18 @@ async def test_get_document_metadata_override_merges_into_document_info(
     assert doc.document_info.metadata["text_length"] == str(len(doc.text))
 
 
+def test_pdf_loader_non_local_filesystem_branch_uses_bytesio(
+    dummy_pdf_path: str,
+) -> None:
+    """Non-\"default\" local FS (``auto_mkdir=True``) reads into BytesIO for pypdf."""
+    from neo4j_graphrag.experimental.components.data_loader import is_default_fs
+
+    fs = LocalFileSystem(auto_mkdir=True)
+    assert is_default_fs(fs) is False
+    text = PdfLoader.load_file(dummy_pdf_path, fs=fs)
+    assert text == "Lorem ipsum dolor sit amet."
+
+
 def test_pdf_loader_backward_compat_reexport_module() -> None:
     """``pdf_loader`` submodule re-exports the same classes as ``data_loader``."""
     from neo4j_graphrag.experimental.components.data_loader import (