深入 GraphRAG 源码：文件摄入解析

微软的 GraphRAG 强大且灵活，越来越多开发者将其引入实际业务场景。但 GraphRAG 的原生功能距离企业级落地仍需大量定制化开发。特别是最新发布的 0.5.0 版本，带来了诸多升级：支持增量更新、Prompt 全面定制、新增 DRIFT 搜索、动态社区选择等特性。然而，伴随数百文件的变更，许多基于旧版（如 0.3.0）修改的代码难以兼容新版本。为此，我们启动了 GraphRAG 源码解析系列文章，帮助大家更高效地完成定制化开发。本篇将聚焦于文件摄取逻辑，为您揭开核心实现的关键细节。

1. 文件摄入

GraphRAG 官方当前仅支持 CSV 和 TXT 格式文件。如果需要支持 PDF、HTML、Markdown 等格式，您需要自行扩展文件读取逻辑。其实，对于 Markdown 文件，只需修改配置即可简单实现支持，稍后会具体演示。

GraphRAG 文件摄取的主入口位于 graphrag/index/input 目录，包含以下关键文件：

- __init__.py
- csv.py
- load_input.py
- text.py

核心函数 load_input 实现文件加载，其定义如下：

async def load_input(
    config: PipelineInputConfig | InputConfig,
    progress_reporter: ProgressReporter | None = None,
    root_dir: str | None = None,
) -> pd.DataFrame:
    """Load the input data for a pipeline."""
    ....
    if config.file_type in loaders:
        ...
        loader = loaders[config.file_type]
        results = await loader(config, progress, storage)
        return cast(pd.DataFrame, results)
...

关键点：通过 loaders 根据文件类型匹配加载器，然后调用对应的读取逻辑。加载器的初始化如下：

loaders: dict[str, Callable[..., Awaitable[pd.DataFrame]]] = {
    text: load_text,
    csv: load_csv,
}

如何扩展支持新格式？

只需在 loaders 中新增文件类型及其加载器，并编写对应实现即可。下面我们通过 txt 格式的实现来具体分析。

2. Txt 文件加载逻辑

以 TXT 格式为例，其加载器load_text定义在text.py中：

async def load(
    config: PipelineInputConfig,
    progress: ProgressReporter | None,
    storage: PipelineStorage,
) -> pd.DataFrame:
    """Load text inputs from a directory."""

    async def load_file(
        path: str, group: dict | None = None, _encoding: str = "utf-8"
    ) -> dict[str, Any]:
        if group is None:
            group = {}
        text = await storage.get(path, encoding="utf-8")
        new_item = {**group, "text": text}
        new_item["id"] = gen_md5_hash(new_item, new_item.keys())
        new_item["title"] = str(Path(path).name)
        return new_item

    files = list(
        storage.find(
            repile(config.file_pattern),
            progress=progress,
            file_filter=config.file_filter,
        )
    )
    ...

    files_loaded = []
    for file, group in files:
        try:
            files_loaded.append(await load_file(file, group))
        except Exception:  # noqa: BLE001 (catching Exception is fine here)
            log.warning("Warning! Error loading file %s. Skipping...", file)
   ...
    return pd.DataFrame(files_loaded)

核心步骤：

1. 根据配置文件中的正则表达式，查找符合条件的文件（默认 .*\\.txt$）。
2. 读取文件内容并生成 MD5 哈希作为唯一标识。
3. 将文件名存为 title 字段，并组织为 DataFrame。

添加 Markdown 支持

如果需要让 GraphRAG 读取 Markdown 文件，只需修改配置文件的 file_pattern：

file_pattern: ".*\\.(txt|md)$"

这将允许 GraphRAG 读取 .md 文件，但要实现基于 Markdown 格式的分段，还需进一步开发 chunk 逻辑。