一个简单易用、便于二次开发的GraphRAG实现

😭 GraphRAG 功能强大且优秀，但其官方实现却难以阅读或修改。

😊 本项目提供了一个更小、更快、更简洁的GraphRAG实现，同时保留了核心功能（详见基准测试和问题讨论）。

🎁 排除tests和提示模板后，nano-graphrag仅约1100行代码。

👌 虽然体积小巧，但它具有可移植性（支持Faiss、Neo4j、Ollama等）、异步支持，并且实现了完全类型化。

如果您正在寻找一种适用于长期用户记忆的多用户RAG解决方案，请查看此项目：memobase :)

安装

从源码安装（推荐）

# 首先克隆本仓库
cd nano-graphrag
pip install -e .

从PyPI安装

pip install nano-graphrag

快速开始

[!TIP]

请在环境变量中设置OpenAI API密钥：export OPENAI_API_KEY="sk-..."。

[!TIP] 如果您使用的是Azure OpenAI API，请参考 .env.example 来配置您的Azure OpenAI信息。然后通过传递 GraphRAG(...,using_azure_openai=True,...) 来启用它。

[!TIP] 如果您使用的是Amazon Bedrock API，请确保已通过 aws configure 等命令正确设置了凭证。之后可以通过如下方式启用：GraphRAG(...,using_amazon_bedrock=True, best_model_id="us.anthropic.claude-3-sonnet-20240229-v1:0", cheap_model_id="us.anthropic.claude-3-haiku-20240307-v1:0",...)。请参阅示例脚本 using_amazon_bedrock.py。

[!TIP]

如果您没有API密钥，可以查看这个示例 no_openai_key_at_all.py，其中使用了transformers和ollama。如果您想使用其他LLM或嵌入模型，请参阅进阶用法。

下载查尔斯·狄更斯的《圣诞颂歌》：

curl https://raw.githubusercontent.com/gusye1234/nano-graphrag/main/tests/mock_data.txt > ./book.txt

然后使用以下Python代码片段：

from nano_graphrag import GraphRAG, QueryParam

graph_func = GraphRAG(working_dir="./dickens")

with open("./book.txt") as f:
    graph_func.insert(f.read())

# 执行全局GraphRAG搜索
print(graph_func.query("这篇故事的主要主题是什么？"))

# 执行局部GraphRAG搜索（我认为这种方式更好且更具扩展性）
print(graph_func.query("这篇故事的主要主题是什么？", param=QueryParam(mode="local")))

下次您从同一working_dir初始化GraphRAG时，它会自动重新加载所有上下文。

批量插入

graph_func.insert(["TEXT1", "TEXT2",...])

with open("./book.txt") as f:
    book = f.read()
    half_len = len(book) // 2
    graph_func.insert(book[:half_len])
    graph_func.insert(book[half_len:])

graph_func = GraphRAG(working_dir="./dickens", enable_naive_rag=True)
...
# 查询
print(rag.query(
      "这篇故事的主要主题是什么？",
      param=QueryParam(mode="naive")
)

异步支持

对于每个方法 NAME(...)，都对应有一个异步方法 aNAME(...)。

await graph_func.ainsert(...)
await graph_func.aquery(...)
...

可用参数

GraphRAG和QueryParam都是Python中的dataclass。您可以使用 help(GraphRAG) 和 help(QueryParam) 查看所有可用参数！或者前往进阶用法部分了解更多选项。

组件

以下是您可以使用的组件：

• 内置 表示我们在 nano-graphrag 中已经实现了该功能。示例 表示我们在 examples 文件夹下的教程中提供了该实现。

• 请查看 examples/benchmarks，以了解各组件之间的对比。

• 我们始终欢迎更多组件的贡献。

进展

# 本地模式
-----报告-----
```csv
id,	content
0,	# 福克斯新闻及媒体和政治领域的重要人物...
1, ...
```
...

# 全局模式
----分析师 3----
重要性得分：100
唐纳德·J·特朗普：经常被讨论与其政治活动相关的内容...
...

from nano_graphrag._op import chunking_by_seperators

GraphRAG(...,chunk_func=chunking_by_seperators,...)

async def my_llm_complete(
    prompt, system_prompt=None, history_messages=[], **kwargs
) -> str:
  # 如果有缓存 KV 数据库，则将其移除
  hashing_kv: BaseKVStorage = kwargs.pop("hashing_kv", None)
  # 剩余的 kwargs 用于调用 LLM，例如 `max_tokens=xxx`
	...
  # 调用您的 LLM
  response = await call_your_LLM(messages, **kwargs)
  return response

# 如需调整最大 token 数量或最大异步请求数量，请进行相应设置
GraphRAG(最佳模型函数=my_llm_complete, 最佳模型最大 token 数量=..., 最佳模型最大异步请求数量=...)
GraphRAG(经济模型函数=my_llm_complete, 经济模型最大 token 数量=..., 经济模型最大异步请求数量=...)

def YOUR_STRING_TO_JSON_FUNC(response: str) -> dict:
  "将字符串响应转换为 JSON"
  ...

@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
async def openai_embedding(texts: list[str]) -> np.ndarray:
    openai_async_client = AsyncOpenAI()
    response = await openai_async_client.embeddings.create(
        model="text-embedding-3-small", input=texts, encoding_format="float"
    )
    return np.array([dp.embedding for dp in response.data])

GraphRAG(embedding_func=your_embed_func, embedding_batch_num=..., embedding_func_max_async=...)

常见问题解答

请查看常见问题解答。

路线图

请参阅ROADMAP.md

贡献

nano-graphrag 欢迎任何形式的贡献。在贡献之前，请阅读这篇文档。

基准测试

• 英文基准测试
• 中文基准测试
• 一项评估笔记本，针对一个多跳 RAG 任务

使用过 nano-graphrag 的项目

• Medical Graph RAG：用于医疗数据的图谱 RAG
• LightRAG：简单快速的检索增强生成
• fast-graphrag：能够智能适应您的用例、数据和查询的 RAG
• HiRAG：基于层次化知识的检索增强生成

如果您的项目使用了 nano-graphrag，欢迎提交 Pull Request，这将有助于其他人信任这个仓库❤️

问题

• nano-graphrag 尚未实现 GraphRAG 的 covariates 功能。
• nano-graphrag 的全局搜索实现方式与原版不同。原版采用类似 MapReduce 的方式将所有社区的信息填充到上下文中，而 nano-graphrag 则只使用前 K 个重要且中心的社区（可通过 QueryParam.global_max_consider_community 控制，默认为 512 个社区）。

nano-graphrag 快速上手指南

nano-graphrag 是一个轻量、快速且易于定制的 GraphRAG（基于图谱的检索增强生成）实现。它保留了核心功能，但代码量仅约 1100 行，支持异步操作和多种后端组件。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：>= 3.9.11
• 前置依赖：
  • 需要配置大模型 API Key（默认支持 OpenAI）。
  • 若使用其他模型（如 Ollama, DeepSeek, Amazon Bedrock），请确保相应服务已就绪。

安装步骤

推荐从源码安装以获取最新功能，也可直接通过 PyPI 安装。

方式一：从源码安装（推荐）

# 克隆仓库
git clone https://github.com/gusye1234/nano-graphrag.git
cd nano-graphrag

# 安装依赖
pip install -e .

方式二：从 PyPI 安装

pip install nano-graphrag

提示：国内用户若下载缓慢，可添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源。

基本使用

1. 配置 API Key

在使用前，请在终端设置环境变量（以 OpenAI 为例）：

export OPENAI_API_KEY="sk-..."

注：若使用 Azure OpenAI 或 Amazon Bedrock，需在代码中传递相应参数（见下文高级用法），或使用本地模型（如 Ollama）则无需此步。

2. 准备测试数据

下载示例文本文件：

curl https://raw.githubusercontent.com/gusye1234/nano-graphrag/main/tests/mock_data.txt > ./book.txt

3. 运行代码

创建一个 Python 文件（例如 main.py），填入以下代码即可体验全局和局部图谱搜索：

from nano_graphrag import GraphRAG, QueryParam

# 初始化 GraphRAG，指定工作目录
graph_func = GraphRAG(working_dir="./dickens")

# 插入数据
with open("./book.txt") as f:
    graph_func.insert(f.read())

# 执行全局图谱搜索 (Global Search)
print(graph_func.query("What are the top themes in this story?"))

# 执行局部图谱搜索 (Local Search，推荐，扩展性更好)
print(graph_func.query("What are the top themes in this story?", param=QueryParam(mode="local")))

说明：

• 首次运行会自动构建知识图谱并保存在 ./dickens 目录。
• 再次初始化相同 working_dir 时，会自动加载已有的上下文，无需重新计算。
• 支持批量插入：graph_func.insert(["TEXT1", "TEXT2", ...])。
• 支持增量插入：重复内容不会重复计算，但社区报告会重新生成。

4. 异步支持

所有同步方法均有对应的异步版本，只需在方法名前加 a 并使用 await：

await graph_func.ainsert(...)
await graph_func.aquery(...)