GraphRAG + Ollama 本地部署配置与源码修改实战

GraphRAG + Ollama 本地部署配置与源码修改实战 | 极客日志

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 # 或 azure_openai_chat
  model: gemma2:9b # 替换为你安装的本地 LLM 模型名称
  model_supports_json: true # 建议开启以支持结构化输出
  max_tokens: 2048
  api_base: http://localhost:11434/v1 # 注意接口路径为 v1
  concurrent_requests: 1
parallelization:
  stagger: 0.3
  async_mode: threaded
embeddings:
  async_mode: threaded
  llm:
    api_key: ollama
    type: openai_embedding # 或 azure_openai_embedding
    model: quentinz/bge-large-zh-v1.5:latest # 替换为你安装的本地 Embedding 模型
    api_base: http://localhost:11434/api # 注意接口路径为 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 库。

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License
"""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:
            # 使用 ollama 库直接调用本地模型
            embedding = ollama.embeddings(model="quentinz/bge-large-zh-v1.5:latest", prompt=inp)
            embedding_list.append(embedding["embedding"])
        return embedding_list
        # 原有 OpenAI 调用逻辑注释如下：

2. 修改 Query 模块中的 Embedding 封装

定位文件：site-packages/graphrag/query/llm/oai/embedding.py 此文件用于查询时的文本向量化，需集成 langchain_community 的 Ollama 封装。

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License
"""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 环境，实现了数据的私有化处理和本地推理。这不仅降低了成本，还提升了数据安全性。开发者可根据实际需求进一步调整模型参数和切片策略，以获得最佳检索效果。

GraphRAG + Ollama 本地部署配置与源码修改实战