利用Milvus向量数据库实现GraphRAG

大家好，GraphRAG技术借助知识图谱，给RAG应用注入了新的动力，使其能够在海量数据中精确检索所需信息。本文将介绍GraphRAG的实现方法，包括如何创建索引以及如何利用Milvus向量数据库进行查询，助力在信息检索的道路上事半功倍。

在运行本文中的代码之前，请确保已安装以下依赖项：

pip install --upgrade pymilvus
pip install git+https://github.com/zc277584121/graphrag.git

通过一个分支仓库来安装GraphRAG，这是因为Milvus的存储功能在本文编写时还未被官方正式合并。

1.数据准备

进行GraphRAG索引需要准备一个文本文件，从Gutenberg项目（https://www.gutenberg.org/）下载一个大约一千行的文本文件，这个文件包含了关于达芬奇的故事。

利用这个数据集，构建一个涉及达芬奇所有关系的知识图谱索引，并使用Milvus向量数据库来检索相关知识，以便回答相关问题。

以下是Python代码，用于下载文本文件并进行初步处理：

import nest_asyncio
nest_asyncio.apply()

import os
import urllib.request

index_root = os.path.join(os.getcwd(), 'graphrag_index')
os.makedirs(os.path.join(index_root, 'input'), exist_ok=True)
url = "https://www.gutenberg.org/cache/epub/7785/pg7785.txt"
file_path = os.path.join(index_root, 'input', 'davinci.txt')
urllib.request.urlretrieve(url, file_path)
with open(file_path, 'r+', encoding='utf-8') as file:
    # 使用文本文件的前934行，因为后面的行与本例无关。
    # 如果想节省API密钥成本，可以截断文本文件以减小大小。
    lines = file.readlines()
    file.seek(0)
    file.writelines(lines[:934])  # 如果想节省API密钥成本，可以减少这个数字。
    file.truncate()

2.配置环境并执行索引流程

现在，使用GraphRAG对文本文件进行索引。首先运行graphrag.index --init命令初始化工作空间。

python -m graphrag.index --init --root ./graphrag_index

在索引的根目录下，能找到一个名为.env的文件。要启用这个文件，请将OpenAI API密钥添加。

注意事项：

• 本例将使用OpenAI模型作为一部分，请准备好自己的API密钥。

• GraphRAG索引的成本相对较高，因为它需要用LLM处理整个文本语料库。运行这个演示可能会花费一些资金。为了节省成本，可以考虑将文本文件缩减尺寸。

运行索引需要一些时间，执行完毕后，可以在./graphrag_index/output/<timestamp>/路径下找到一个新创建的文件夹，里面包含了多个parquet格式的文件。

执行以下命令开始索引过程：

python -m graphrag.index --root ./graphrag_index

3.使用Milvus向量数据库进行查询

在查询阶段，使用Milvus来存储GraphRAG本地搜索中实体描述的向量嵌入。

这种方法将知识图谱的结构化数据与输入文档的非结构化数据相结合，为LLM提供了额外的相关实体信息，从而能够得出更准确的答案。

import os

import pandas as pd
import tiktoken
from graphrag.query.con