← ブログに戻る

Knowledge GraphをPython500行でスクラッチ実装 GraphRAG動くまで

/ 6分で読めます / knowledge-graph · graphrag · python · networkx · llm

Neo4jを立てず、LangChainも入れず、標準ライブラリ+networkxだけで500行のPythonを書いたら、GraphRAGの最小構成が動きました。動いた瞬間の私の第一声は「え、これでいいの?」でした。先に結論から言います。GraphRAGは概念としては巨大に見えますが、最小構成の部品は4つしかありません。エンティティ抽出、グラフ構築、コミュニティ検出、要約付き検索。この4つをPythonで素直に繋ぐだけなら、Neo4jもLangChainも要りません。むしろ最初はそれらを入れないほうが、GraphRAGの中で本当に難しい部分がどこなのかを掴めます。

Knowledge Graphを500行で組む4ステージのパイプライン

対象読者と、あえて外したもの

対象は「GraphRAGを触ってみたいけど、Neo4j Auraの請求書が怖い人」です。私も同じでした。運用フェーズになったらもちろんNeo4jやNeptuneに移せばよくて、この記事の500行はその前段の「本当に自社データでGraphRAGが効くのか」を1日で確かめるためのものです。

あえて外したもの:

  • Neo4j / Nebula / Neptune などのグラフDB
  • LangChain / LlamaIndex などのフレームワーク
  • ベクトルDB(FAISSやChromaは近いけど、今回はコサイン類似度をnumpyで手書き)
  • LangGraphのGraph(これは実行フローのグラフであって、ナレッジグラフではない。ここは混同されがち)

外した理由は同じで、「500行の中に何を書いたか」を最後まで自分で把握したかったからです。フレームワークを噛ませると、動いた瞬間に「何が効いてるのか」が分からなくなる。

私が過去に3回踏んだ罠です。

LangGraphのGraphとKnowledge Graphの違い

先に地雷を1つだけ潰しておきます。

LangGraphのStateGraphは、エージェントの実行フローを有向グラフで表現するランタイムです。ナレッジグラフは、テキストから抽出したエンティティと関係を保持する知識表現です。両方に「グラフ」がついていて紛らわしいですが、目的も持ち回すデータもまったく別物です。LangGraphのノードはPython関数、エッジは状態遷移条件。ナレッジグラフのノードはエンティティ(人、組織、概念)、エッジはそれらの関係(所属、参照、原因など)。この記事で言う「Knowledge Graph」は後者、つまりMicrosoft Researchが2024年に発表したGraphRAGが構築するタイプのグラフです。

500行の内訳(先出し)

ネタバレしておくと、500行のうち意味のあるロジックはこう分かれます。

モジュール行数役割
extractor.py約120行LLMでエンティティ・関係抽出
graph.py約80行networkx.MultiDiGraphのラッパー
community.py約90行Leiden相当のクラスタリング(python-louvain使用)
summarizer.py約60行コミュニティ要約生成
retriever.py約110行クエリからコミュニティ検索+回答
cli.py約40行動作確認用CLI

合計500行前後。標準ライブラリ以外はopenai、networkx、python-louvain、numpyの4つだけ。

Step 1: エンティティ・関係抽出(120行)

一番泥臭くて、一番効くステップです。テキストチャンクをLLMに投げて、エンティティのリストと関係のリストをJSONで返してもらいます。

import json
from openai import OpenAI

client = OpenAI()

EXTRACT_PROMPT = """
以下のテキストからエンティティと関係を抽出してJSONで返してください。

エンティティタイプ: PERSON, ORG, CONCEPT, TECH, PRODUCT
関係タイプ: developed_by, part_of, references, competes_with, uses

出力フォーマット:
{
  "entities": [{"id": "...", "type": "...", "name": "..."}],
  "relations": [{"src": "...", "rel": "...", "dst": "..."}]
}

テキスト:
{text}
"""

def extract_entities_relations(text: str) -> dict:
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": EXTRACT_PROMPT.format(text=text)}],
        response_format={"type": "json_object"},
        temperature=0.1,
    )
    return json.loads(resp.choices[0].message.content)

ここでハマりやすいのが、エンティティのID正規化です。同じ「Neo4j」というエンティティが「neo4j」「Neo 4j」「Neo4J社」と表記揺れするだけで、グラフが分断されます。私は最初これに気づかず、100ノードのグラフを組んだつもりで実際は50個の孤立サブグラフでした。正規化は素朴に小文字化+空白除去でも、最初は十分です。凝りたくなる気持ちは分かりますが、動くまでは我慢。

Step 2: networkxでグラフ構築(80行)

グラフDBを使わない代わりにnetworkx.MultiDiGraphを使います。理由は多重エッジ(同じノード対に複数の関係)を許すためです。ナレッジグラフでは「AがBを開発した」と「AがBを買収した」が並列で存在します。

import networkx as nx

class KnowledgeGraph:
    def __init__(self):
        self.g = nx.MultiDiGraph()

    def add_from_extraction(self, extraction: dict):
        for e in extraction["entities"]:
            self.g.add_node(e["id"], type=e["type"], name=e["name"])
        for r in extraction["relations"]:
            self.g.add_edge(r["src"], r["dst"], key=r["rel"], rel=r["rel"])

    def neighbors_of(self, node_id: str, depth: int = 1) -> set:
        seen = {node_id}
        frontier = {node_id}
        for _ in range(depth):
            nxt = set()
            for n in frontier:
                nxt.update(self.g.successors(n))
                nxt.update(self.g.predecessors(n))
            frontier = nxt - seen
            seen |= frontier
        return seen

networkxはスケール限界があります。ノード10万を超えるとneighbors_ofが急に遅くなります。ですが「1日で効果検証」の用途では十分。むしろこの制約が「まず小さく試す」という健全な圧力になります。

Step 3: コミュニティ検出(90行)

GraphRAGの肝は「密に繋がったノード群」をコミュニティとして抽出し、そのコミュニティ単位で要約を作ることです。オリジナルはLeidenアルゴリズムですが、Python標準品質で近い挙動をするのはpython-louvainです。Louvainのほうが古典的ですが、500行スケールなら差は気にならないレベルです。

import community as louvain

def detect_communities(kg: KnowledgeGraph) -> dict:
    undirected = kg.g.to_undirected()
    partition = louvain.best_partition(undirected)
    communities = {}
    for node, cid in partition.items():
        communities.setdefault(cid, []).append(node)
    return communities

コミュニティ数はデータ次第です。私が試した100ページ相当のテキストでは、大体8〜15コミュニティに落ち着きました。これが200になると要約コストが跳ねるので、大きいテキストに対してはresolutionパラメータで粒度を調整します。

Step 4: コミュニティ要約(60行)

各コミュニティに含まれるエンティティと関係を、LLMに1段落で要約させます。

SUMMARY_PROMPT = """
以下のエンティティ群と関係を1段落で要約してください。
このコミュニティが何についてのものかを、200字以内で。

エンティティ:
{entities}

関係:
{relations}
"""

def summarize_community(kg, community_nodes: list) -> str:
    subgraph = kg.g.subgraph(community_nodes)
    entities = [kg.g.nodes[n]["name"] for n in subgraph.nodes]
    relations = [f"{u} -{d['rel']}-> {v}" for u, v, d in subgraph.edges(data=True)]
    prompt = SUMMARY_PROMPT.format(
        entities="\n".join(entities),
        relations="\n".join(relations),
    )
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,
    )
    return resp.choices[0].message.content

この要約が、後段の検索インデックスになります。

ここまでで、テキスト100ページから10個くらいの「コミュニティ要約」が手に入りました。要は「この文書群の中でまとまった話題の目次」ができた状態です。

Step 5: 検索と回答(110行)

ユーザーのクエリに対して、コミュニティ要約をベクトル類似度で並べ替え、上位を回答生成に使います。ベクトルDBは入れず、numpyでコサイン類似度を計算します。

import numpy as np

def embed(text: str) -> np.ndarray:
    resp = client.embeddings.create(model="text-embedding-3-small", input=text)
    return np.array(resp.data[0].embedding)

def retrieve(query: str, community_summaries: dict, top_k: int = 3) -> list:
    q = embed(query)
    scored = []
    for cid, summary in community_summaries.items():
        s = embed(summary)
        score = np.dot(q, s) / (np.linalg.norm(q) * np.linalg.norm(s))
        scored.append((score, cid, summary))
    scored.sort(reverse=True)
    return scored[:top_k]

def answer(query: str, community_summaries: dict) -> str:
    top = retrieve(query, community_summaries)
    context = "\n---\n".join(s for _, _, s in top)
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": f"以下の文脈から質問に答えてください。\n\n文脈:\n{context}\n\n質問: {query}",
        }],
    )
    return resp.choices[0].message.content

100ページ程度の埋め込みならnumpyで十分な速度で動きます。1000ページを超えたらFAISSに移す判断ラインです。

動かして分かったこと

500行が動いてから2週間、社内Wiki相当のテキスト120ページで検証しました。得た知見は3つ。

1つ目、GraphRAGが強いのは「横断質問」だけ。 「XプロジェクトとYプロジェクトに共通する課題は?」のような質問では、ベクトルRAGより明確に良い回答が出ました。一方で「Xプロジェクトの担当者は誰?」のような単一事実検索は、素朴なベクトル検索の圧勝でした。GraphRAGはハンマー、ベクトルRAGはドライバー。使い所を間違えると両方悲しくなります。

2つ目、コストはStep 1が9割。 LLMを使うのは抽出・要約・回答生成の3箇所ですが、圧倒的にStep 1のエンティティ抽出のトークン消費が多いです。私の120ページで、抽出だけで約$4かかりました。要約は$0.3、回答は1クエリ数セント。テキストが増えるとStep 1が線形に効いてくるので、最初からgpt-4o-miniのような安い抽出専用モデルを選ぶのが正解です。

3つ目、Neo4jは本当に必要になってから入れる。 networkxの限界は明確に来ますが、「来てから」の判断で十分です。私は最初「どうせスケールするならNeo4j Auraから始めよう」と3回考えて、3回とも仮説検証フェーズで頓挫させました。500行で動く実装があると、その仮説検証を1日で終わらせられます。

まとめ

  • GraphRAGの最小構成は4部品(抽出・グラフ構築・コミュニティ・要約検索)
  • Neo4jもLangChainも要らず、networkx+python-louvainで500行に収まる
  • Step 1のエンティティ抽出だけがコストと精度の両方を握っている
  • スケール限界は明確に来るが、「来てから」で判断できる

GraphRAGを本格的に自社データに適用する前段として、500行のミニマム実装で「効くのか効かないのか」を1日で確かめる。これができるだけで、Neo4j Auraの月額を試算する前の判断材料が全部揃います。

関連

関連する記事

記事一覧を見る →