GraphRAG + Ollama 本地部署配置与源码修改实战
背景与目的
微软开源的 GraphRAG(Graph-based Retrieval-Augmented Generation)是一个强大的知识图谱构建工具,但其官方版本默认强依赖于 OpenAI 的闭源大模型 API。这导致在数据隐私敏感、网络受限或成本控制严格的场景下,直接部署面临较大困难。为了打破这一限制,本文通过修改 GraphRAG 的源码及配置文件,使其支持接入 Ollama 提供的本地 Embedding 模型和大语言模型(LLM),从而实现完全本地化的知识图谱构建与检索。
一、环境准备与安装
1. 基础环境要求
确保系统已安装 Python 3.10 至 3.12 版本。建议使用虚拟环境管理依赖。
2. 安装 GraphRAG
通过 pip 安装 GraphRAG 核心包:
pip install graphrag
3. 安装并配置 Ollama
访问 Ollama 官网下载并安装服务端。安装完成后,拉取所需的本地模型。
推荐使用的模型组合如下:
- Embedding 模型:
quentinz/bge-large-zh-v1.5:latest (支持中文语义理解)
- LLM 大模型:
gemma2:9b (轻量级且性能较好的开源模型)
拉取命令示例:
ollama pull quentinz/bge-large-zh-v1.5
ollama pull gemma2:9b
二、初始化工作区
- 创建用于存储知识数据的文件夹,目前 GraphRAG 主要支持
.txt 和 .csv 格式文本。
- 将待处理的数据文件放入
/ragtest/input 目录下。
- 初始化工作区,运行以下命令生成配置文件:
graphrag init --root ragtest
执行后会在 ragtest 目录下生成 .env 和 settings.yaml 两个关键文件。
.env: 包含环境变量,如 API_KEY,需根据实际配置修改。
settings.yaml: 包含管道运行的详细设置,是本次修改的核心。
三、修改配置文件支持本地模型
打开 settings.yaml,重点修改 llm 和 embeddings 部分,将 API 类型指向本地 Ollama 服务。
encoding_model: cl100k_base
skip_workflows: []
llm:
api_key: ollama
type: openai_chat
model: gemma2:9b
model_supports_json: true
max_tokens: 2048
api_base: http://localhost:11434/v1
concurrent_requests: 1
parallelization:
stagger: 0.3
async_mode: threaded
embeddings:
async_mode: threaded
llm:
api_key: ollama
type: openai_embedding
model: quentinz/bge-large-zh-v1.5:latest
api_base: http://localhost:11434/api
concurrent_requests: 1
chunks:
size: 300
overlap: 100
group_by_columns: [id]
input:
type: file
file_type: text
base_dir: "input"
file_encoding: utf-8
file_pattern: ".*\\.txt$"
cache:
type: file
base_dir: "cache"
storage:
type: file
base_dir: "output/${timestamp}/artifacts"
reporting:
type: file
base_dir: "output/${timestamp}/reports"
entity_extraction:
prompt: "prompts/entity_extraction.txt"
entity_types: [organization,person,geo,event]
max_gleanings: 0
summarize_descriptions:
prompt: "prompts/summarize_descriptions.txt"
max_length: 500
claim_extraction:
prompt: "prompts/claim_extraction.txt"
description: "Any claims or facts that could be relevant to information discovery."
max_gleanings: 0
community_report:
prompt: "prompts/community_report.txt"
max_length: 2000
max_input_length: 8000
cluster_graph:
max_cluster_size: 10
embed_graph:
enabled: false
umap:
enabled: false
snapshots:
graphml: false
raw_entities: false
top_level_nodes: false
local_search:
max_tokens: 5000
global_search:
max_tokens: 5000
四、修改源码适配本地调用
仅修改配置文件有时不足以完全兼容,部分底层调用仍需修改源码以适配 Ollama 的 SDK 方式。
1. 修改 Embedding 生成器
定位文件:site-packages/graphrag/llm/openai/openai_embeddings_llm.py
该文件负责生成文本向量。我们需要将其从调用 OpenAI 客户端改为调用 ollama 库。
"""The EmbeddingsLLM class."""
from typing_extensions import Unpack
from graphrag.llm.base import BaseLLM
from graphrag.llm.types import (
EmbeddingInput,
EmbeddingOutput,
LLMInput,
)
from .openai_configuration import OpenAIConfiguration
from .types import OpenAIClientTypes
import ollama
class OpenAIEmbeddingsLLM(BaseLLM[EmbeddingInput, EmbeddingOutput]):
"""A text-embedding generator LLM."""
_client: OpenAIClientTypes
_configuration: OpenAIConfiguration
def __init__(self, client: OpenAIClientTypes, configuration: OpenAIConfiguration):
self.client = client
self.configuration = configuration
async def _execute_llm(
self, input: EmbeddingInput, **kwargs: Unpack[LLMInput]
) -> EmbeddingOutput | None:
args = {
"model": self.configuration.model,
**(kwargs.get("model_parameters") or {}),
}
embedding_list = []
for inp in input:
embedding = ollama.embeddings(model="quentinz/bge-large-zh-v1.5:latest", prompt=inp)
embedding_list.append(embedding["embedding"])
return embedding_list
2. 修改 Query 模块中的 Embedding 封装
定位文件:site-packages/graphrag/query/llm/oai/embedding.py
此文件用于查询时的文本向量化,需集成 langchain_community 的 Ollama 封装。
"""OpenAI Embedding model implementation."""
import asyncio
from collections.abc import Callable
from typing import Any
import numpy as np
import tiktoken
from tenacity import (
AsyncRetrying,
RetryError,
Retrying,
retry_if_exception_type,
stop_after_attempt,
wait_exponential_jitter,
)
from graphrag.query.llm.base import BaseTextEmbedding
from graphrag.query.llm.oai.base import OpenAILLMImpl
from graphrag.query.llm.oai.typing import (
OPENAI_RETRY_ERROR_TYPES,
OpenaiApiType,
)
from graphrag.query.llm.text_utils import chunk_text
from graphrag.query.progress import StatusReporter
from langchain_community.embeddings import OllamaEmbeddings
class OpenAIEmbedding(BaseTextEmbedding, OpenAILLMImpl):
"""Wrapper for OpenAI Embedding models."""
def __init__(
self,
api_key: str | None = None,
azure_ad_token_provider: Callable | None = None,
model: str = "text-embedding-3-small",
deployment_name: str | None = None,
api_base: | = ,
api_version: | = ,
api_type: OpenaiApiType = OpenaiApiType.OpenAI,
organization: | = ,
encoding_name: = ,
max_tokens: = ,
max_retries: = ,
request_timeout: = ,
retry_error_types: [[BaseException]] = OPENAI_RETRY_ERROR_TYPES,
reporter: StatusReporter | = ,
):
OpenAILLMImpl.__init__(
=,
api_key=api_key,
azure_ad_token_provider=azure_ad_token_provider,
deployment_name=deployment_name,
api_base=api_base,
api_version=api_version,
api_type=api_type,
organization=organization,
max_retries=max_retries,
request_timeout=request_timeout,
reporter=reporter,
)
.model = model
.encoding_name = encoding_name
.max_tokens = max_tokens
.token_encoder = tiktoken.get_encoding(.encoding_name)
.retry_error_types = retry_error_types
() -> []:
token_chunks = chunk_text(
text=text, token_encoder=.token_encoder, max_tokens=.max_tokens
)
chunk_embeddings = []
chunk_lens = []
chunk token_chunks:
:
embedding, chunk_len = ._embed_with_retry(chunk, **kwargs)
chunk_embeddings.append(embedding)
chunk_lens.append(chunk_len)
Exception e:
._reporter.error(
message=,
details={.__class__.__name__: (e)},
)
chunk_embeddings = np.average(chunk_embeddings, axis=, weights=chunk_lens)
chunk_embeddings = chunk_embeddings / np.linalg.norm(chunk_embeddings)
chunk_embeddings.tolist()
() -> []:
token_chunks = chunk_text(
text=text, token_encoder=.token_encoder, max_tokens=.max_tokens
)
chunk_embeddings = []
chunk_lens = []
embedding_results = asyncio.gather(*[
._aembed_with_retry(chunk, **kwargs) chunk token_chunks
])
embedding_results = [result result embedding_results result[]]
chunk_embeddings = [result[] result embedding_results]
chunk_lens = [result[] result embedding_results]
chunk_embeddings = np.average(chunk_embeddings, axis=, weights=chunk_lens)
chunk_embeddings = chunk_embeddings / np.linalg.norm(chunk_embeddings)
chunk_embeddings.tolist()
() -> [[], ]:
:
retryer = Retrying(
stop=stop_after_attempt(.max_retries),
wait=wait_exponential_jitter(=),
reraise=,
retry=retry_if_exception_type(.retry_error_types),
)
attempt retryer:
attempt:
embedding = (
OllamaEmbeddings(
model=.model,
).embed_query(text)
[]
)
(embedding, (text))
RetryError e:
._reporter.error(
message=,
details={.__class__.__name__: (e)},
)
([], )
:
([], )
() -> [[], ]:
:
retryer = AsyncRetrying(
stop=stop_after_attempt(.max_retries),
wait=wait_exponential_jitter(=),
reraise=,
retry=retry_if_exception_type(.retry_error_types),
)
attempt retryer:
attempt:
embedding = (
OllamaEmbeddings(
model=.model,
).embed_query(text) [] )
(embedding, (text))
RetryError e:
._reporter.error(
message=,
details={.__class__.__name__: (e)},
)
([], )
:
([], )
五、构建索引与效果测试
1. 构建知识图谱索引
完成上述配置和代码修改后,运行 GraphRAG 构建命令:
graphrag index --root ragtest
构建过程会经历实体提取、关系发现、社区摘要等阶段,耗时取决于数据量和硬件性能。
2. 查询测试
构建完成后,可进行两种模式的查询测试:
- Local Search: 基于局部上下文进行回答,适合具体细节查询。
- Global Search: 基于全局知识图谱进行回答,适合宏观总结。
六、常见问题排查
- API Key 错误: 确保
settings.yaml 中 api_key 设置为 ollama,而非真实的 OpenAI Key。
- 端口冲突: 确认 Ollama 服务运行在
http://localhost:11434,且未被防火墙拦截。
- 模型未找到: 检查
settings.yaml 中的模型名称是否与 ollama list 输出的名称完全一致。
- 内存不足: 本地部署大模型对显存和内存要求较高,若遇到 OOM 错误,请尝试减小
max_tokens 或使用更小参数的模型(如 gemma2:2b)。
结语
通过上述步骤,我们成功将 GraphRAG 迁移至本地 Ollama 环境,实现了数据的私有化处理和本地推理。这不仅降低了成本,还提升了数据安全性。开发者可根据实际需求进一步调整模型参数和切片策略,以获得最佳检索效果。