Gemini API で構築する GraphRAG — 知識グラフ × ベクトル検索の本番ハイブリッド RAG

# pip install google-genai neo4j
import os
import json
from google import genai
from google.genai import types
 
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
EXTRACT_GRAPH_TOOL = types.Tool(
    function_declarations=[
        types.FunctionDeclaration(
            name="store_knowledge_graph",
            description="文書から抽出した知識グラフの三つ組を保存する",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "triples": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(
                            type=types.Type.OBJECT,
                            properties={
                                "subject": types.Schema(type=types.Type.STRING, description="主語エンティティ名"),
                                "subject_type": types.Schema(type=types.Type.STRING, description="Person / Product / Tech / Concept など"),
                                "predicate": types.Schema(type=types.Type.STRING, description="関係名（uses, depends_on, created_by など）"),
                                "object": types.Schema(type=types.Type.STRING, description="目的語エンティティ名"),
                                "object_type": types.Schema(type=types.Type.STRING),
                                "evidence": types.Schema(type=types.Type.STRING, description="根拠となる原文の引用"),
                            },
                            required=["subject", "subject_type", "predicate", "object", "object_type", "evidence"],
                        ),
                    )
                },
                required=["triples"],
            ),
        )
    ]
)
 
def extract_triples(document: str, source_id: str) -> list[dict]:
    """文書から知識グラフの三つ組を抽出する。
    返り値は [{subject, subject_type, predicate, object, object_type, evidence, source_id}, ...]
    """
    prompt = f"""次の文書を読み、登場するエンティティ間の関係を三つ組として抽出してください。
- 同じエンティティは同じ表記で統一すること（例: Gemini API と「Gemini の API」は Gemini API に正規化）
- 推測ではなく文書に明記された関係のみを抽出すること
- evidence には根拠となる原文を1〜2文で引用すること
 
文書:
{document}
"""
    response = client.models.generate_content(
        model="gemini-2.5-pro",
        contents=prompt,
        config=types.GenerateContentConfig(
            tools=[EXTRACT_GRAPH_TOOL],
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode="ANY")
            ),
            temperature=0.0,
        ),
    )
 
    # Function call から triples を取り出す
    fc = response.candidates[0].content.parts[0].function_call
    if fc is None or fc.name != "store_knowledge_graph":
        return []
    triples = fc.args.get("triples", [])
    for t in triples:
        t["source_id"] = source_id
    return triples
 
if __name__ == "__main__":
    sample = """Dolice Labs の Claude Lab は Cloudflare Workers にデプロイされている。
Claude Lab は Stripe を使って課金を行い、コンテンツは MDX で書かれている。
Claude Lab を作成したのは政樹である。"""
    triples = extract_triples(sample, source_id="claudelab-readme")
    print(json.dumps(triples, ensure_ascii=False, indent=2))
    # 期待出力:
    # [
    #   {"subject": "Claude Lab", "predicate": "deployed_on", "object": "Cloudflare Workers", ...},
    #   {"subject": "Claude Lab", "predicate": "uses", "object": "Stripe", ...},
    #   {"subject": "Claude Lab", "predicate": "created_by", "object": "政樹", ...},
    #   ...
    # ]

# pip install neo4j
from neo4j import GraphDatabase
 
driver = GraphDatabase.driver(
    os.environ["NEO4J_URI"],
    auth=(os.environ["NEO4J_USER"], os.environ["NEO4J_PASSWORD"]),
)
 
UPSERT_QUERY = """
MERGE (s:Entity {name: $subject})
  ON CREATE SET s.type = $subject_type
MERGE (o:Entity {name: $object})
  ON CREATE SET o.type = $object_type
MERGE (s)-[r:RELATES {predicate: $predicate, source_id: $source_id}]->(o)
  ON CREATE SET r.evidence = $evidence, r.created_at = timestamp()
"""
 
def write_triples_to_neo4j(triples: list[dict]) -> None:
    """三つ組を Neo4j に冪等に書き込む。"""
    with driver.session() as session:
        for t in triples:
            try:
                session.run(UPSERT_QUERY, **t)
            except Exception as e:
                # 1件失敗しても残りを書き込めるようにロギングして継続
                print(f"⚠️ failed to upsert triple: {t} — {e}")
 
def index_document(document: str, source_id: str, vector_store) -> None:
    """文書をグラフとベクトルストアの両方に書き込む。"""
    triples = extract_triples(document, source_id=source_id)
    if not triples:
        print(f"⚠️ no triples extracted for {source_id}")
        return
    write_triples_to_neo4j(triples)
 
    # 同じ文書をベクトルストアにも書き込む
    embedding = client.models.embed_content(
        model="gemini-embedding-001",
        contents=document,
        config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT"),
    ).embeddings[0].values
    vector_store.upsert(
        id=source_id,
        vector=embedding,
        metadata={"text": document, "triples_count": len(triples)},
    )

ROUTE_TOOL = types.Tool(
    function_declarations=[
        types.FunctionDeclaration(
            name="route_question",
            description="質問を最適な検索ルートに振り分ける",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "route": types.Schema(
                        type=types.Type.STRING,
                        enum=["GRAPH_ONLY", "VECTOR_ONLY", "HYBRID"],
                        description="GRAPH_ONLY=エンティティ間の明確な関係を問う質問 / VECTOR_ONLY=曖昧で文脈依存の質問 / HYBRID=両方が必要",
                    ),
                    "entities": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(type=types.Type.STRING),
                        description="質問に登場するエンティティ名（正規化済み）",
                    ),
                },
                required=["route", "entities"],
            ),
        )
    ]
)
 
def route_query(question: str) -> dict:
    """質問を分類してルーティング情報を返す。"""
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=f"次の質問を分類してください: {question}",
        config=types.GenerateContentConfig(
            tools=[ROUTE_TOOL],
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode="ANY")
            ),
            temperature=0.0,
        ),
    )
    return response.candidates[0].content.parts[0].function_call.args
 
def graph_traverse(entities: list[str], hops: int = 2) -> list[dict]:
    """指定エンティティから N ホップ以内のサブグラフを取得する。"""
    cypher = """
    MATCH path = (start:Entity)-[*1..$hops]-(neighbor:Entity)
    WHERE start.name IN $entities
    RETURN start.name AS start, neighbor.name AS neighbor,
           [rel IN relationships(path) | {predicate: rel.predicate, evidence: rel.evidence, source_id: rel.source_id}] AS path
    LIMIT 50
    """
    with driver.session() as session:
        result = session.run(cypher, entities=entities, hops=hops)
        return [r.data() for r in result]
 
def vector_search(question: str, top_k: int = 8, vector_store=None) -> list[dict]:
    """ベクトル検索でチャンクを取得する。"""
    query_embedding = client.models.embed_content(
        model="gemini-embedding-001",
        contents=question,
        config=types.EmbedContentConfig(task_type="RETRIEVAL_QUERY"),
    ).embeddings[0].values
    return vector_store.query(vector=query_embedding, top_k=top_k, include_metadata=True)

def subgraph_to_text(subgraph: list[dict]) -> str:
    """サブグラフを Gemini が読みやすいテキスト形式に変換する。"""
    lines = []
    for row in subgraph:
        path_str = " → ".join(
            f"[{p['predicate']}] (出典: {p['source_id']})"
            for p in row["path"]
        )
        lines.append(f"{row['start']} {path_str} {row['neighbor']}")
    return "\n".join(lines)
 
def answer_with_graphrag(question: str, vector_store) -> str:
    """GraphRAG で質問に回答する。"""
    routing = route_query(question)
    route = routing["route"]
    entities = routing["entities"]
 
    graph_context = ""
    vector_context = ""
 
    if route in ("GRAPH_ONLY", "HYBRID") and entities:
        subgraph = graph_traverse(entities, hops=2)
        graph_context = subgraph_to_text(subgraph)
 
    if route in ("VECTOR_ONLY", "HYBRID"):
        chunks = vector_search(question, top_k=8, vector_store=vector_store)
        vector_context = "\n---\n".join(c["metadata"]["text"] for c in chunks)
 
    prompt = f"""次のコンテキストを参考に、質問に答えてください。回答には必ず出典の source_id を [角括弧] で示してください。
推測で補わず、コンテキストに無い情報は『不明』と答えてください。
 
[知識グラフ（関係性ベース）]
{graph_context or 'なし'}
 
[文書チャンク（意味類似ベース）]
{vector_context or 'なし'}
 
[質問]
{question}
"""
    response = client.models.generate_content(
        model="gemini-2.5-pro",
        contents=prompt,
        config=types.GenerateContentConfig(temperature=0.2),
    )
    return response.text