neo4jライブラリを使ったハイブリッドGraphRAG

はじめに

昨日のこの投稿の続きで、ハイブリッドな検索を使ったGraphRAGシステムの構築をまとめた。

LangChainのLCELのようにパイプ「｜」を繋げていくイメージで考えていたが、想像していたものとは違った。以下で確かめて欲しい。

情報源

Hybrid Retrieval for GraphRAG Applications Using the Neo4j GraphRAG Package for Python この記事のコードを試した。

実行結果

neo4jデータベースに接続、OpenAIのAPIキーを環境変数に設定の部分は、この投稿と同じである。

HybridRetrieverを用意

この投稿のVectorRetrieverを次のようにHybridRetrieverに置き換える。

# Hybrid Retriever

from neo4j import GraphDatabase
from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
from neo4j_graphrag.retrievers import HybridRetriever

embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
retriever = HybridRetriever(
    driver,
    vector_index_name="moviePlotsEmbedding",
    fulltext_index_name="movieFulltext",
    embedder=embedder,
    return_properties=["title", "plot"],
)

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 is the name of the movie set in 1375 in Imperial China?"
response = rag.search(query_text=query_text, retriever_config={"top_k": 3})
print(response.answer)

実行結果

The name of the movie set in 1375 in Imperial China is "Musa the Warrior (Musa)."

まとめ

neo4j_graphragパッケージにHybridRetrieverクラスが用意されているので、そのオブジェクトをretrieverとしてGraphRAGクラスを呼べば、簡単に実現できる。