chore: import upstream snapshot with attribution

packages/graphrag-input/graphrag_input/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""GraphRAG input document loading package."""
+
+from graphrag_input.get_property import get_property
+from graphrag_input.input_config import InputConfig
+from graphrag_input.input_reader import InputReader
+from graphrag_input.input_reader_factory import create_input_reader
+from graphrag_input.input_type import InputType
+from graphrag_input.text_document import TextDocument
+
+__all__ = [
+    "InputConfig",
+    "InputReader",
+    "InputType",
+    "TextDocument",
+    "create_input_reader",
+    "get_property",
+]

packages/graphrag-input/graphrag_input/csv.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'CSVFileReader' model."""
+
+import csv
+import io
+import logging
+import sys
+
+from graphrag_input.structured_file_reader import StructuredFileReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+try:
+    csv.field_size_limit(sys.maxsize)
+except OverflowError:
+    csv.field_size_limit(100 * 1024 * 1024)
+
+
+class CSVFileReader(StructuredFileReader):
+    """Reader implementation for csv files."""
+
+    def __init__(self, file_pattern: str | None = None, **kwargs):
+        super().__init__(
+            file_pattern=file_pattern if file_pattern is not None else ".*\\.csv$",
+            **kwargs,
+        )
+
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a csv file into a list of documents.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - list with a TextDocument for each row in the file.
+        """
+        file = await self._storage.get(path, encoding=self._encoding)
+
+        reader = csv.DictReader(io.StringIO(file))
+        rows = list(reader)
+        return await self.process_data_columns(rows, path)

packages/graphrag-input/graphrag_input/get_property.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Utility for retrieving properties from nested dictionaries."""
+
+from typing import Any
+
+
+def get_property(data: dict[str, Any], path: str) -> Any:
+    """Retrieve a property from a dictionary using dot notation.
+
+    Parameters
+    ----------
+    data : dict[str, Any]
+        The dictionary to retrieve the property from.
+    path : str
+        A dot-separated string representing the path to the property (e.g., "foo.bar.baz").
+
+    Returns
+    -------
+    Any
+        The value at the specified path.
+
+    Raises
+    ------
+    KeyError
+        If the path does not exist in the dictionary.
+    """
+    keys = path.split(".")
+    current = data
+    for key in keys:
+        if not isinstance(current, dict) or key not in current:
+            msg = f"Property '{path}' not found"
+            raise KeyError(msg)
+        current = current[key]
+    return current

packages/graphrag-input/graphrag_input/hashing.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Hashing utilities."""
+
+from collections.abc import Iterable
+from hashlib import sha512
+from typing import Any
+
+
+def gen_sha512_hash(item: dict[str, Any], hashcode: Iterable[str]) -> str:
+    """Generate a SHA512 hash.
+
+    Parameters
+    ----------
+    item : dict[str, Any]
+        The dictionary containing values to hash.
+    hashcode : Iterable[str]
+        The keys to include in the hash.
+
+    Returns
+    -------
+    str
+        The SHA512 hash as a hexadecimal string.
+    """
+    hashed = "".join([str(item[column]) for column in hashcode])
+    return f"{sha512(hashed.encode('utf-8'), usedforsecurity=False).hexdigest()}"

packages/graphrag-input/graphrag_input/input_config.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Parameterization settings for the default configuration."""
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from graphrag_input.input_type import InputType
+
+
+class InputConfig(BaseModel):
+    """The default configuration section for Input."""
+
+    model_config = ConfigDict(extra="allow")
+    """Allow extra fields to support custom reader implementations."""
+
+    type: str = Field(
+        description="The input file type to use.",
+        default=InputType.Text,
+    )
+    encoding: str | None = Field(
+        description="The input file encoding to use.",
+        default=None,
+    )
+    file_pattern: str | None = Field(
+        description="The input file pattern to use.",
+        default=None,
+    )
+    id_column: str | None = Field(
+        description="The input ID column to use.",
+        default=None,
+    )
+    title_column: str | None = Field(
+        description="The input title column to use.",
+        default=None,
+    )
+    text_column: str | None = Field(
+        description="The input text column to use.",
+        default=None,
+    )

packages/graphrag-input/graphrag_input/input_reader.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'InputReader' model."""
+
+from __future__ import annotations
+
+import logging
+import re
+from abc import ABCMeta, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from graphrag_storage import Storage
+
+    from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class InputReader(metaclass=ABCMeta):
+    """Provide a cache interface for the pipeline."""
+
+    def __init__(
+        self,
+        storage: Storage,
+        file_pattern: str,
+        encoding: str = "utf-8",
+        **kwargs,
+    ):
+        self._storage = storage
+        self._encoding = encoding
+        self._file_pattern = file_pattern
+
+    async def read_files(self) -> list[TextDocument]:
+        """Load all files from storage and return them as a single list."""
+        return [doc async for doc in self]
+
+    def __aiter__(self) -> AsyncIterator[TextDocument]:
+        """Return the async iterator, enabling `async for doc in reader`."""
+        return self._iterate_files()
+
+    async def _iterate_files(self) -> AsyncIterator[TextDocument]:
+        """Async generator that yields documents one at a time as files are loaded."""
+        files = list(self._storage.find(re.compile(self._file_pattern)))
+        if len(files) == 0:
+            msg = f"No {self._file_pattern} matches found in storage"
+            logger.warning(msg)
+            return
+
+        file_count = len(files)
+        doc_count = 0
+
+        for file in files:
+            try:
+                for doc in await self.read_file(file):
+                    doc_count += 1
+                    yield doc
+            except Exception as e:  # noqa: BLE001 (catching Exception is fine here)
+                logger.warning("Warning! Error loading file %s. Skipping...", file)
+                logger.warning("Error: %s", e)
+
+        logger.info(
+            "Found %d %s files, loading %d",
+            file_count,
+            self._file_pattern,
+            doc_count,
+        )
+        logger.info(
+            "Total number of unfiltered %s rows: %d",
+            self._file_pattern,
+            doc_count,
+        )
+
+    @abstractmethod
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a file into a list of documents.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - List with an entry for each document in the file.
+        """

packages/graphrag-input/graphrag_input/input_reader_factory.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'InputReaderFactory' model."""
+
+import logging
+from collections.abc import Callable
+
+from graphrag_common.factory import Factory
+from graphrag_common.factory.factory import ServiceScope
+from graphrag_storage.storage import Storage
+
+from graphrag_input.input_config import InputConfig
+from graphrag_input.input_reader import InputReader
+from graphrag_input.input_type import InputType
+
+logger = logging.getLogger(__name__)
+
+
+class InputReaderFactory(Factory[InputReader]):
+    """Factory for creating Input Reader instances."""
+
+
+input_reader_factory = InputReaderFactory()
+
+
+def register_input_reader(
+    input_reader_type: str,
+    input_reader_initializer: Callable[..., InputReader],
+    scope: ServiceScope = "transient",
+) -> None:
+    """Register a custom input reader implementation.
+
+    Args
+    ----
+        - input_reader_type: str
+            The input reader id to register.
+        - input_reader_initializer: Callable[..., InputReader]
+            The input reader initializer to register.
+    """
+    input_reader_factory.register(input_reader_type, input_reader_initializer, scope)
+
+
+def create_input_reader(config: InputConfig, storage: Storage) -> InputReader:
+    """Create an input reader implementation based on the given configuration.
+
+    Args
+    ----
+        - config: InputConfig
+            The input reader configuration to use.
+        - storage: Storage | None
+            The storage implementation to use for reading the files.
+
+    Returns
+    -------
+        InputReader
+            The created input reader implementation.
+    """
+    config_model = config.model_dump()
+    input_strategy = config.type
+
+    if input_strategy not in input_reader_factory:
+        match input_strategy:
+            case InputType.Csv:
+                from graphrag_input.csv import CSVFileReader
+
+                register_input_reader(InputType.Csv, CSVFileReader)
+            case InputType.Text:
+                from graphrag_input.text import TextFileReader
+
+                register_input_reader(InputType.Text, TextFileReader)
+            case InputType.Json:
+                from graphrag_input.json import JSONFileReader
+
+                register_input_reader(InputType.Json, JSONFileReader)
+            case InputType.JsonLines:
+                from graphrag_input.jsonl import JSONLinesFileReader
+
+                register_input_reader(InputType.JsonLines, JSONLinesFileReader)
+            case InputType.MarkItDown:
+                from graphrag_input.markitdown import MarkItDownFileReader
+
+                register_input_reader(InputType.MarkItDown, MarkItDownFileReader)
+            case InputType.Parquet:
+                from graphrag_input.parquet import ParquetFileReader
+
+                register_input_reader(InputType.Parquet, ParquetFileReader)
+            case _:
+                msg = f"InputConfig.type '{input_strategy}' is not registered in the InputReaderFactory. Registered types: {', '.join(input_reader_factory.keys())}."
+                raise ValueError(msg)
+
+    config_model["storage"] = storage
+
+    return input_reader_factory.create(input_strategy, init_args=config_model)

packages/graphrag-input/graphrag_input/input_type.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing input file type enum."""
+
+from enum import StrEnum
+
+
+class InputType(StrEnum):
+    """The input file type for the pipeline."""
+
+    Csv = "csv"
+    """The CSV input type."""
+    Text = "text"
+    """The text input type."""
+    Json = "json"
+    """The JSON input type."""
+    JsonLines = "jsonl"
+    """The JSON Lines input type."""
+    MarkItDown = "markitdown"
+    """The MarkItDown input type."""
+    Parquet = "parquet"
+    """The Parquet input type."""
+
+    def __repr__(self):
+        """Get a string representation."""
+        return f'"{self.value}"'

packages/graphrag-input/graphrag_input/json.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'JSONFileReader' model."""
+
+import json
+import logging
+
+from graphrag_input.structured_file_reader import StructuredFileReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class JSONFileReader(StructuredFileReader):
+    """Reader implementation for json files."""
+
+    def __init__(self, file_pattern: str | None = None, **kwargs):
+        super().__init__(
+            file_pattern=file_pattern if file_pattern is not None else ".*\\.json$",
+            **kwargs,
+        )
+
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a JSON file into a list of documents.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - list with a TextDocument for each row in the file.
+        """
+        text = await self._storage.get(path, encoding=self._encoding)
+        as_json = json.loads(text)
+        # json file could just be a single object, or an array of objects
+        rows = as_json if isinstance(as_json, list) else [as_json]
+        return await self.process_data_columns(rows, path)

packages/graphrag-input/graphrag_input/jsonl.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'JSONLinesFileReader' model."""
+
+import json
+import logging
+
+from graphrag_input.structured_file_reader import StructuredFileReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class JSONLinesFileReader(StructuredFileReader):
+    """Reader implementation for json lines files."""
+
+    def __init__(self, file_pattern: str | None = None, **kwargs):
+        super().__init__(
+            file_pattern=file_pattern if file_pattern is not None else ".*\\.jsonl$",
+            **kwargs,
+        )
+
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a JSON lines file into a list of documents.
+
+        This differs from standard JSON files in that each line is a separate JSON object.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - list with a TextDocument for each row in the file.
+        """
+        text = await self._storage.get(path, encoding=self._encoding)
+        rows = [json.loads(line) for line in text.splitlines()]
+        return await self.process_data_columns(rows, path)

packages/graphrag-input/graphrag_input/markitdown.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'TextFileReader' model."""
+
+import logging
+from io import BytesIO
+from pathlib import Path
+
+from markitdown import MarkItDown, StreamInfo
+
+from graphrag_input.hashing import gen_sha512_hash
+from graphrag_input.input_reader import InputReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class MarkItDownFileReader(InputReader):
+    """Reader implementation for any file type supported by markitdown.
+
+    https://github.com/microsoft/markitdown
+    """
+
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a text file into a DataFrame of documents.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - list with a TextDocument for each row in the file.
+        """
+        bytes = await self._storage.get(path, encoding=self._encoding, as_bytes=True)
+        md = MarkItDown()
+        result = md.convert_stream(
+            BytesIO(bytes), stream_info=StreamInfo(extension=Path(path).suffix)
+        )
+        text = result.markdown
+
+        document = TextDocument(
+            id=gen_sha512_hash({"text": text}, ["text"]),
+            title=result.title if result.title else str(Path(path).name),
+            text=text,
+            creation_date=await self._storage.get_creation_date(path),
+            raw_data=None,
+        )
+        return [document]

packages/graphrag-input/graphrag_input/parquet.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'ParquetFileReader' model."""
+
+import io
+import logging
+
+import pyarrow.parquet as pq
+
+from graphrag_input.structured_file_reader import StructuredFileReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class ParquetFileReader(StructuredFileReader):
+    """Reader implementation for parquet files."""
+
+    def __init__(self, file_pattern: str | None = None, **kwargs):
+        super().__init__(
+            file_pattern=file_pattern if file_pattern is not None else ".*\\.parquet$",
+            **kwargs,
+        )
+
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a parquet file into a list of documents.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - list with a TextDocument for each row in the file.
+        """
+        file_bytes = await self._storage.get(path, as_bytes=True)
+        table = pq.read_table(io.BytesIO(file_bytes))
+        rows = table.to_pylist()
+        return await self.process_data_columns(rows, path)

packages/graphrag-input/graphrag_input/structured_file_reader.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'StructuredFileReader' model."""
+
+import logging
+from typing import Any
+
+from graphrag_input.get_property import get_property
+from graphrag_input.hashing import gen_sha512_hash
+from graphrag_input.input_reader import InputReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class StructuredFileReader(InputReader):
+    """Base reader implementation for structured files such as csv and json."""
+
+    def __init__(
+        self,
+        id_column: str | None = None,
+        title_column: str | None = None,
+        text_column: str = "text",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._id_column = id_column
+        self._title_column = title_column
+        self._text_column = text_column
+
+    async def process_data_columns(
+        self,
+        rows: list[dict[str, Any]],
+        path: str,
+    ) -> list[TextDocument]:
+        """Process configured data columns from a list of loaded dicts."""
+        documents = []
+        for index, row in enumerate(rows):
+            # text column is required - harvest from dict
+            text = get_property(row, self._text_column)
+            # id is optional - harvest from dict or hash from text
+            id = (
+                get_property(row, self._id_column)
+                if self._id_column
+                else gen_sha512_hash({"text": text}, ["text"])
+            )
+            # title is optional - harvest from dict or use filename
+            num = f" ({index})" if len(rows) > 1 else ""
+            title = (
+                get_property(row, self._title_column)
+                if self._title_column
+                else f"{path}{num}"
+            )
+            creation_date = await self._storage.get_creation_date(path)
+            documents.append(
+                TextDocument(
+                    id=id,
+                    title=title,
+                    text=text,
+                    creation_date=creation_date,
+                    raw_data=row,
+                )
+            )
+        return documents

packages/graphrag-input/graphrag_input/text.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing 'TextFileReader' model."""
+
+import logging
+from pathlib import Path
+
+from graphrag_input.hashing import gen_sha512_hash
+from graphrag_input.input_reader import InputReader
+from graphrag_input.text_document import TextDocument
+
+logger = logging.getLogger(__name__)
+
+
+class TextFileReader(InputReader):
+    """Reader implementation for text files."""
+
+    def __init__(self, file_pattern: str | None = None, **kwargs):
+        super().__init__(
+            file_pattern=file_pattern if file_pattern is not None else ".*\\.txt$",
+            **kwargs,
+        )
+
+    async def read_file(self, path: str) -> list[TextDocument]:
+        """Read a text file into a list of documents.
+
+        Args:
+            - path - The path to read the file from.
+
+        Returns
+        -------
+            - output - list with a TextDocument for each row in the file.
+        """
+        text = await self._storage.get(path, encoding=self._encoding)
+        document = TextDocument(
+            id=gen_sha512_hash({"text": text}, ["text"]),
+            title=str(Path(path).name),
+            text=text,
+            creation_date=await self._storage.get_creation_date(path),
+            raw_data=None,
+        )
+        return [document]

packages/graphrag-input/graphrag_input/text_document.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""TextDocument dataclass."""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+from graphrag_input.get_property import get_property
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TextDocument:
+    """The TextDocument holds relevant content for GraphRAG indexing."""
+
+    id: str
+    """Unique identifier for the document."""
+    text: str
+    """The main text content of the document."""
+    title: str
+    """The title of the document."""
+    creation_date: str
+    """The creation date of the document, ISO-8601 format."""
+    raw_data: dict[str, Any] | None = None
+    """Raw data from source document."""
+
+    def get(self, field: str, default_value: Any = None) -> Any:
+        """
+        Get a single field from the TextDocument.
+
+        Functions like the get method on a dictionary, returning default_value if the field is not found.
+
+        Supports nested fields using dot notation.
+
+        This takes a two step approach for flexibility:
+        1. If the field is one of the standard text document fields (id, title, text, creation_date), just grab it directly. This accommodates unstructured text for example, which just has the standard fields.
+        2. Otherwise. try to extract it from the raw_data dict. This allows users to specify any column from the original input file.
+
+        """
+        if field in ["id", "title", "text", "creation_date"]:
+            return getattr(self, field)
+
+        raw = self.raw_data or {}
+        try:
+            return get_property(raw, field)
+        except KeyError:
+            return default_value
+
+    def collect(self, fields: list[str]) -> dict[str, Any]:
+        """Extract data fields from a TextDocument into a dict."""
+        data = {}
+        for field in fields:
+            value = self.get(field)
+            if value is not None:
+                data[field] = value
+        return data