Remove fs

Author: ali-sedaghatbaf

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

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

src/neo4j_graphrag/experimental/components/data_loader.py

@@ -16,15 +16,11 @@
 
 from __future__ import annotations
 
-import io
 from abc import abstractmethod
 from pathlib import Path
-from typing import Dict, Optional, Union, cast
+from typing import Dict, Optional, Union
 
-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 (
@@ -35,10 +31,6 @@
 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.
@@ -54,23 +46,18 @@ 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,
-        fs: AbstractFileSystem,
-    ) -> str:
+    def load_file(file: str) -> str:
         """Parse a PDF file and return extracted text."""
         try:
-            with fs.open(file, "rb") as fp:
-                stream = fp if is_default_fs(fs) else io.BytesIO(fp.read())
-                pdf = pypdf.PdfReader(stream)
+            with open(file, "rb") as fp:
+                pdf = pypdf.PdfReader(fp)
                 num_pages = len(pdf.pages)
                 text_parts = (
                     pdf.pages[page].extract_text() for page in range(num_pages)
@@ -83,15 +70,10 @@ 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)
-        if isinstance(fs, str):
-            fs = fsspec.filesystem(fs)
-        elif fs is None:
-            fs = LocalFileSystem()
-        text = self.load_file(filepath, fs)
+        text = self.load_file(filepath)
         return LoadedDocument(
             text=text,
             document_info=DocumentInfo(
@@ -106,30 +88,22 @@ class MarkdownLoader(DataLoader):
     """Loads UTF-8 Markdown (``.md`` / ``.markdown``) files as plain text."""
 
     @staticmethod
-    def load_file(
-        file: str,
-        fs: AbstractFileSystem,
-    ) -> str:
+    def load_file(file: str) -> str:
         try:
-            with fs.open(file, "rb") as fp:
+            with open(file, "rb") as fp:
                 raw = fp.read()
-            return cast(str, raw.decode("utf-8"))
+            return 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)
-        if isinstance(fs, str):
-            fs = fsspec.filesystem(fs)
-        elif fs is None:
-            fs = LocalFileSystem()
-        text = MarkdownLoader.load_file(filepath, fs)
+        text = MarkdownLoader.load_file(filepath)
         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,16 +83,13 @@ 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, fs=fs)
+            return await PdfLoader().run(filepath=path_str, metadata=metadata)
         if suffix in (".md", ".markdown"):
-            return await MarkdownLoader().run(
-                filepath=path_str, metadata=metadata, fs=fs
-            )
+            return await MarkdownLoader().run(filepath=path_str, metadata=metadata)
         raise UnsupportedDocumentFormatError(
             f"Unsupported document format: {suffix!r}. "
             f"Supported: .pdf, .md, .markdown"

tests/unit/experimental/components/test_data_loader.py

@@ -20,8 +20,6 @@
 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,
@@ -49,31 +47,25 @@ 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, fs=LocalFileSystem())
+    actual_content = pdf_loader.load_file(dummy_pdf_path)
     assert actual_content == expected_content
 
 
 def test_pdf_processing_error(pdf_loader: PdfLoader, dummy_pdf_path: str) -> None:
-    with patch(
-        "fsspec.implementations.local.LocalFileSystem.open",
-        side_effect=Exception("Failed to open"),
-    ):
+    with patch("builtins.open", side_effect=Exception("Failed to open")):
         with pytest.raises(PdfLoaderError):
-            pdf_loader.load_file(dummy_pdf_path, fs=LocalFileSystem())
+            pdf_loader.load_file(dummy_pdf_path)
 
 
 def test_markdown_processing_error(dummy_md_path: str) -> None:
-    with patch(
-        "fsspec.implementations.local.LocalFileSystem.open",
-        side_effect=Exception("Failed to open"),
-    ):
+    with patch("builtins.open", side_effect=Exception("Failed to open")):
         with pytest.raises(MarkdownLoadError):
-            MarkdownLoader.load_file(dummy_md_path, fs=LocalFileSystem())
+            MarkdownLoader.load_file(dummy_md_path)
 
 
 def test_markdown_loading() -> None:
     md_path = str(BASE_DIR / "sample_data/hello.md")
-    text = MarkdownLoader.load_file(md_path, fs=LocalFileSystem())
+    text = MarkdownLoader.load_file(md_path)
     assert "# Hello" in text
     assert "Markdown **content**" in text
 
@@ -89,7 +81,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` (default ``fs``)."""
+    """``PdfLoader.run`` wraps ``load_file`` with :class:`DocumentInfo`."""
     pdf_path = BASE_DIR / "sample_data/lorem_ipsum.pdf"
     loader = PdfLoader()
     doc = await loader.run(filepath=pdf_path)
@@ -98,25 +90,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_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()
@@ -132,9 +105,8 @@ 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, fs=fs)
+        return await super().run(filepath=filepath, metadata=metadata)
 
     def get_document_metadata(
         self, text: str, metadata: dict[str, str] | None = None
@@ -158,18 +130,6 @@ 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 (