neo4jライブラリを使ったGraphRAG

はじめに

この投稿では、neo4jを用いてナレッジグラフを作成し、ナレッジグラフを外部情報とするRAGシステムを構築した。

この投稿では、neo4jライブラリを使って、簡単にRAGシステムを構築する方法を試したので、その内容をまとめた。

情報源

Getting Started With the Neo4j GraphRAG Python Package この記事のコードを試した。記事のコードを修正した箇所があるので、そこは本文中でコメントする。
Enriching Vector Search With Graph Traversal Using the GraphRAG Python Package この記事のコードも試した。上記とほぼ同じ内容。
Hybrid Retrieval for GraphRAG Applications Using the Neo4j GraphRAG Package for Python この記事のコードは未だ試していない。今後試してみる予定。

実行結果

用意されているneo4jデータベースに接続

# Establish a connection to your Neo4j database

from neo4j import GraphDatabase
# Demo database credentials
URI = "neo4j+s://demo.neo4jlabs.com"
AUTH = ("recommendations", "recommendations")
# Connect to Neo4j database
driver = GraphDatabase.driver(URI, auth=AUTH)

OpenAIのAPIキーを環境変数に設定

# Set OpenAI API key

import os

os.environ['OPENAI_API_KEY'] = 'sk-*****'

GraphRAGのretrieverを用意

# Prepare GraphRAG's retriever

from neo4j_graphrag.retrievers import VectorRetriever
from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings

embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
retriever = VectorRetriever(
    driver,
    index_name="moviePlotsEmbedding",
    embedder=embedder,
    return_properties=["title", "plot"],
)

パッケージ名を「neo4j-graphrag」から「neo4j_graphrag」に変更した。

GraphRAGでパイプラインを構築

# GraphRAG

from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.generation import GraphRAG

# LLM
# Note: the OPENAI_API_KEY must be in the env vars
llm = OpenAILLM(model_name="gpt-4o-mini", model_params={"temperature": 0})

# Initialize the RAG pipline
rag = GraphRAG(retriever=retriever, llm=llm)

# Query the graph
query_text = "What movies are sad rommances?"
response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
print(response.answer)

実行結果

Based on the context provided, the following movies can be considered sad romances:

1. **Autumn in New York** - This film is a romantic drama about an aging playboy who falls for a sweet, but terminally ill, young woman, which adds a poignant and sad element to the romance.

2. **Bed of Roses** - While it has romantic elements, the story of a young career girl being swept off her feet by a shy florist may also carry emotional weight, depending on the character's journeys.

These films explore themes of love intertwined with sadness and emotional challenges.

まとめ

neo4jのgraphragパッケージを使用すると、非常に簡単にナレッジグラフを外部情報とするRAGシステムが構築できることがわかった。

次は、情報源3.のGraphRAGでハイブリッド検索を行うRAGシステムを試してみようと思っている。