在我的上一篇文章中，我详细讨论了GraphRAG的两种查询模式，包括各自的使用场景和运行原理。然而，在实际操作中，我发现GraphRAG在执行Query时，使用了类似于pandas的方法来将本地目录的parquet文件读入内存。这引发了一个问题：如果parquet文件过大，那么内存可能无法承受其压力。更糟糕的是，如果你使用docker部署GraphRAG应用，那么容器可能会报OOM（内存溢出）错误。
那么，如何解决这个问题呢？一种可行的方案是，我们可以避免使用parquet文件作为知识图谱的存储方式，而是直接利用一个生产级别的图数据库来存储数据。通过这种方式，我们可以直接从图数据库中查询图数据，从而避免发生OOM错误。
因此，在接下来的文章中，我计划基于Neo4j来实现自定义的GraphRAG查询。由于我目前在工作中的主要技术栈是langchain，所以我将采用langchain来进行讲解和示例展示。
在阅读这篇文章之前，你需要先了解GraphRAG自带的查询流程和原理，更多详细参考我之前写的一篇文章喂饭教程！全网首发对小白友好的GraphRAG查询流程全揭秘，这里为了内容的完整性，我只简单带过。
官方查询使用
GraphRAG提供了一个命令行接口(CLI)工具，这使得用户可以在已经构建好的知识图谱上进行查询：
poetry run poe query
--root ./msgraphrag
--method local
--community_level 2
--response_type "Multiple Paragraphs"
'路飞的船员有哪些？'
以下是几个重要的查询参数说明：
community_level：这是社区的层级参数，在使用leiden算法识别社区时会生成。这个数值的高低代表了社区的大小；值越高，社区越小。默认值是2。你可以使用下面的命令来分析在Neo4j数据库中导入的不同级别的社区数量：
MATCH (c:__Community__)
RETURN c.level AS level, COUNT(*) AS count
ORDER BY level
response_type：这是一个描述响应类型和格式的参数。例如，如果你需要多段落形式的输出，那么你可以设置为默认的"Multiple Paragraphs"。如果你需要单句形式的回答，你可以将其设置为“Single Sentence”。method：这个参数表示Query模式，可选项为local（本地）或global（全局）。local模式常用于处理针对具体事实的问题，而global模式则支持基于高层次语义理解的概括性问题，这是适应于QFS型查询任务。
原理
Local
在local模式查询中，我们主要结合知识图谱中的结构化信息和原始文档中的非结构化数据，以此来构建增强生成的上下文。然后，我们借助于语言模型（LLM）来获取响应。因此，local模式特别适合用于回答关于具体事实的问题，例如查询某个实体的特定信息或者探索实体间的关系等。大致流程如下：
根据query和对话历史，借助实体的description_embedding向量进行相似度检索从知识图谱中识别出最相关的实体
找到这些实体更多的信息，包含：
关联的原始文本块。提取其文本内容 关联的社区。提取其社区报告 关联的实体。提取其实体描述信息 关联的关系。提取其关系描述信息 关联的协变量。由于默认不生成，这里忽略
对这些提取的信息进行排序与筛选，最终形成参考的上下文
借助LLM和prompt，生成最终响应
Global
global模式查询的架构相对简单，它采用了分布式计算中的Map-Reduce架构。可以简单概括为：
Map: 查询指定community_level的所有社区报告，使用Query、conversation_history 和 这些社区报告 调用多次LLM生成多个带有评分的中间响应 Reduce：对中间响应进行排序后，将这些中间响应给LLM作为参考，让LLM输出最终答案
global查询模式的问题是响应质量会受到输入的社区级别参数的影响。如果层次较低，则报告较为详细，响应可能会更全面，但所需的时间和模型成本较高。所以具体需要在使用时根据实际情况作权衡考虑。当然，你也可以在产品层将这个社区参数交给用户自己去设定。
自定义GraphRAG查询在自定义查询之前，我们首先需要把数据导入到Neo4j中，我在这篇文章中也提到了对的方法，具体不再赘述。参考上面提到的原理，我将使用Langchain在Neo4j中的知识图谱基础上自定义实现一个local模式的查询。我本地使用的docker启动的neo4j：
docker run \
-p 7474:7474 -p 7687:7687 \
-e NEO4J_apoc_export_file_enabled=true \
-e NEO4J_apoc_import_file_enabled=true \
-e NEO4J_apoc_import_file_use__neo4j__config=true \
-e NEO4J_PLUGINS=\[\"apoc\"\] \
-d \
neo4j:latest
1. 创建数据库索引加速查询
由于我们需要使用query进行相似实体的搜索，而这个操作主要依赖于实体的description_embedding属性来完成。因此，我们首先需要对这个属性建立索引。建立索引的目的有两个：一是提高查询速度，二是为这个字段增加相似度检索的功能：
CREATE VECTOR INDEX entity_index IF NOT EXISTS FOR (e:__Entity__) ON e.description_embedding
OPTIONS {indexConfig: {
`vector.dimensions`: 1536,
`vector.similarity_function`: 'cosine'
}}
这个Neo4j的Cypher查询用于创建一个名为entity_index的向量索引，如果该索引已经存在，则不再创建。它是在节点标签为Entity的节点上，特别是在description_embedding属性上创建的。
选项中的indexConfig设置了一些有关此索引的配置信息：
vector.dimensions: 定义了向量的维度，即每个向量包含的元素数量。在这个例子中，每个description_embedding向量都应该有1536个元素。我用的openai的embedding，其向量维度是1536，你可以根据自己的数据调整这个值vector.similarity_function: 定义了计算两个向量相似性的函数。在这个例子中，我们用的是余弦相似性 ('cosine')，这是一种常用于处理高维数据的相似性度量方式。
2. 验证检索
接下来我们使用Langchain的Neo4jVector验证这个索引是否生效：
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores.neo4j_vector import Neo4jVector
from pprint import pprint
entity_vector = Neo4jVector.from_existing_index(
embedding=OpenAIEmbeddings(),
url="bolt://localhost:7687",
username="neo4j",
password="1234qwer",
index_name='entity_index',
text_node_property='description',
)
result = entity_vector.similarity_search("路飞的船员有哪些", k=20)
pprint(result)
输出结果如下，说明一切正确：
有一点需要注意，text_node_property='description' 这个参数必不可少，这是用来指定节点的text属性字段，用于查询出结果之后从description字段获取值用来填充Langchain对象Document的page_content属性。
3. 找到这些实体更多的信息
阅读Langchain源码，发现执行entity_vector.similarity_search会调用similarity_search_with_score_by_vector方法：
def similarity_search_with_score_by_vector(
self, embedding: List[float], k: int = 4, **kwargs: Any
) -> List[Tuple[Document, float]]:
default_retrieval = (
f"RETURN node.`{self.text_node_property}` AS text, score, "
f"node {{.*, `{self.text_node_property}`: Null, "
f"`{self.embedding_node_property}`: Null, id: Null }} AS metadata"
)
retrieval_query = (
self.retrieval_query if self.retrieval_query else default_retrieval
)
read_query = _get_search_index_query(self.search_type) + retrieval_query
parameters = {
"index": self.index_name,
"k": k,
"embedding": embedding,
"keyword_index": self.keyword_index_name,
"query": kwargs["query"],
}
results = self.query(read_query, params=parameters)
docs = [
(
Document(
page_content=result["text"],
metadata={
k: v for k, v in result["metadata"].items() if v is not None
},
),
result["score"],
)
for result in results
]
return docs
我发现在发送到Neo4j的查询read_query中，存在一个变量：retrieval_query。这个变量可以被用来进一步检索更多的信息，我们来定义这样一个变量，用于从neo4j获取节点关联的其他信息：
retrieval_query_template = """
//接收向量检索输出的node，在此基础上进一步检索
WITH collect(node) as nodes
//查找最相关的文本块，输出text属性
WITH
collect {
UNWIND nodes as n
MATCH (n)<-[:HAS_ENTITY]->(c:__Chunk__)
WITH c, count(distinct n) as freq
RETURN c.text AS chunkText
ORDER BY freq DESC
LIMIT {{ topChunks }}
} AS text_mapping,
//查找最相关的社区，输出summary摘要（如果没有weight，用cypher设定）
collect {
UNWIND nodes as n
MATCH (n)-[:IN_COMMUNITY]->(c:__Community__)
WITH c, c.rank as rank, c.weight AS weight
RETURN c.summary
ORDER BY rank, weight DESC
LIMIT {{ topCommunities }}
} AS report_mapping,
//查找最相关的其他实体（nodes外部），输出描述
collect {
UNWIND nodes as n
MATCH (n)-[r:RELATED]-(m)
WHERE NOT m IN nodes
RETURN r.description AS descriptionText
ORDER BY r.rank, r.weight DESC
LIMIT {{ topOutsideRels }}
} as outsideRels,
//查找最相关的其他实体（nodes内部），输出描述
collect {
UNWIND nodes as n
MATCH (n)-[r:RELATED]-(m)
WHERE m IN nodes
RETURN r.description AS descriptionText
ORDER BY r.rank, r.weight DESC
LIMIT {{ topInsideRels }}
} as insideRels,
//输出这些实体本身的描述
collect {
UNWIND nodes as n
RETURN n.description AS descriptionText
} as entities
//返回text，score，metadata三个字段
RETURN "-----Chunks-----" + apoc.text.join(text_mapping, ' | ') +
"\n-----Reports-----" + apoc.text.join(report_mapping,' | ') +
"\nRelationships-----" + apoc.text.join(outsideRels + insideRels, ' | ') +
"\nEntities-----" + apoc.text.join(entities, " | ") as text, 1.0 AS score, {source:''} AS metadata
"""
这里的Cypher语句虽然较长，但其实并不复杂。就是对向量检索输出的node搜集后，检索更多相关信息（社区报告、其他实体、关系等），最后合并输出
注意这里必须输出text，score，metadata三个属性，这是Langchain构建Document的需要。
每个部分的候选词（文本单元、社区报告等）数量是可配置的。因为GraphRAG实现基于token的过滤比较复杂一些，我在此就先简化处理了，这里给出的默认配置如下:
topChunks = 33
topCommunities = 3
topOutsideRels = 10
topInsideRels = 10
topEntities = 10
好了，准备工作完成之后我们再来测试一下neo4j的查询
retrieval_query = Template(retrieval_query_template).render(**{
"topChunks": topChunks,
"topCommunities": topCommunities,
"topOutsideRels": topOutsideRels,
"topInsideRels": topInsideRels,
})
entity_vector = Neo4jVector.from_existing_index(
embedding=OpenAIEmbeddings(),
url="bolt://localhost:7687",
username="neo4j",
password="1234qwer",
index_name='entity_index',
retrieval_query=retrieval_query
)
result = entity_vector.similarity_search("路飞的船员有哪些", k=topEntities)
pprint(result[0].page_content)
4. 调用LLM
从neo4j中搜索到所有相关数据之后，我们就可以利用Langchain的chain来进行问答了，这里我选用便宜的gpt-4o-mini模型和QA类型的chain来做问答：
llm = ChatOpenAI(model="gpt-4o-mini")
chain = RetrievalQAWithSourcesChain.from_chain_type(
llm,
chain_type="stuff",
retriever=entity_vector.as_retriever(search_kwargs={'k': topEntities}),
verbose=True
)
response = chain.invoke(
{"question": "路飞的船员有哪些？"},
return_only_outputs=True,
)
print(response['answer'])
返回结果正常, 只是相比GraphRAG的响应没有那么丰富，你可以微调prompt来满足你的需求：
> Entering new RetrievalQAWithSourcesChain chain...
> Finished chain.
路飞的船员包括：
1. **蒙奇·D·路飞**（船长）
2. **罗罗诺亚·索隆**（战斗员）
3. **娜美**（航海士）
4. **乌索普**（狙击手）
5. **山治**（厨师）
6. **托尼托尼·乔巴**（船医）
7. **妮可·罗宾**（考古学家）
8. **弗兰奇**（船匠）
9. **布鲁克**（音乐家）
10. **甚平**（舵手）
这些成员共同组成了草帽海贼团，追求成为海贼王的梦想。
好了，到此为止我们已经完成了基于neo4j的检索和Langchain的查询了，有效解决了GraphRAG进行Search会占用大量内存的问题，是不是也挺简单。
总结如果你习惯使用LlamaIndex框架，也可以采用类似方法实现。当然在实际使用中，你还可以根据自身的需要，进一步优化这里的检索召回策略。甚至可以结合查询重写、其他索引（如普通向量索引）策略、Rerank模型等实现更复杂的RAG范式，以获得最佳效果，这很好地扩充了Microsft GraphRAG的应用场景。
除了local模式的查询外，global模式也可自定义实现。区别在于local模式的关键在于召回相关上下文，global模式关键在于map-reduce过程，感兴趣的朋友可以自己实现，如果有疑问可以点赞关注并私信我，我会帮忙给与回答。
值得进一步说明的是，上述基于Neo4j的查询和问答系统并不完美，直接应用于生产环境还存在许多问题。在我的实际AIGC工作中，我为了满足完整产品的需求，还做了大量其他的相关工作。
