Neo4jでベクトル検索する方法

目次

0_前置き
1_ベクトルインデックスの作成
2_ノードの作成
3_ベクトル検索の実行
4_まとめ

0_前置き

Neo4jでベクトル検索ができるようになりました。

Neo4jは言わずと知れたグラフデータベースです。ノードとリレーションシップにより、ナレッジグラフを作成することができます。

ChatGPTの登場で、ベクトルデータを効率的に保存するベクトルデータベースの利用が話題になりました。ベクトルデータベースを使えば、ベクトルの類似度等を用いて値を検索することができます。

そしてそんな中、2023年8月頃に、Neo4jもベクトル検索に対応しました。

Neo4jでベクトル検索を利用する

Neo4jでは、ノードのプロパティにベクトルの値を格納することで、データベース内の特定のラベルに対して、ベクトル検索を行うことができます。

公式の記事に記されたコードでは、productsラベルを持つノードから、ベクトル検索を用いて5件のノードを検出しています。

さらに、4行目以降のクエリでは、この検索結果を用いて、グラフデータベースのリレーションシップに紐づく情報を検出しています。
（ここでは、見つかった製品のカタログとブランド名、製品を購入した上で☆5を付けたカスタマーのレビュー）

// find products by similarity search in vector index
CALL db.index.vector.queryNodes('products', 5, $embedding) yield node as product, score

// enrich with additional explicit relationships from the knowledge graph
MATCH (product)-[:HAS_CATEGORY]->(cat), (product)-[:BY_BRAND]->(brand)
OPTIONAL MATCH (product)<-[:BOUGHT]-(customer)-[rated:RATED]->(product) 
         WHERE rated.rating = 5
OPTIONAL MATCH (product)-[:HAS_REVIEW]->(review)<-[:WROTE]-(customer) 

RETURN product.Name, product.Description, brand.Name, cat.Name, 
       collect(review { .Date, .Text })[0..5] as reviews

このように、Neo4jはグラフデータベースとしての利点を保ちながら、ベクトル検索にも対応することで、検索に柔軟なクエリを用いることができるようになったわけです。

ただし、まだ日本語での情報が少ないみたいでした。
ですので、Neo4jでのベクトル検索の使い方を以下にまとめました。
気になる方はぜひ使ってみてください。

Neo4jの導入

とは言え、まだほぼ最新のNeo4j のバージョン5.13でベータ版になったくらいの機能なので、古いバージョンだと使えません。
クラウドバージョンのNeo4j auraなら、無料で使えて、最新バージョンに管理されているので、さくっと試したいなら、Neo4j auraを使うのが良いでしょう。

Neo4j auraの使い方（本筋じゃないのでさくっと説明）

① Neo4j auraに登録する。

② インスタンスを立ち上げる。（1個まで無料）

③ pythonから使う場合は、neo4jのドライバを入れて、クエリを送信する。
（ブラウザから直接クエリを入力することもできる。）

④ ドライバの初期化を行うコード

from neo4j import GraphDatabase

uri = os.environ["NEO4J_URI"]
username = "neo4j"
password = os.environ["NEO4J_PASSWORD"]
driver = GraphDatabase.driver(uri, auth=(username, password))

⑤ ドライバの初期化を行い、クエリを送信する。

with driver.session() as session:
        result = session.run("クエリ")

⑥　結果の編集
クエリの返り値が単一のレコードなら、result.single()で値を得る。
複数行のレコードならば、resultをfor文で値を得る。
詳しくは、chatGPTに聞けば分かる。

Neo4jでベクトル検索

記事のメインです。ベクトル検索の方法を説明します。
以下は、公式のドキュメントです。

Vector search index - Cypher Manual

1_ベクトルインデックスの作成

まずNeo4jのベクトルインデックスを作成します。
クエリは以下の通りです。

CALL db.index.vector.createNodeIndex(
    'Message', 'Message', 'embedding', 1536, 'cosine')

ここで、各引数は以下の通りです。

1. インデックス名　(何でも構いません）
2. ラベル名
3. 'embedding'
4. ベクトルの次元数（ここではopenAIのembeddingを使用）
5. ベクトルの比較手法

def create_index():
    """メッセージラベルのベクターインデックスを作成する。"""
    with driver.session() as session:
        session.run(
            """
            CALL db.index.vector.createNodeIndex(
                $label, $label, 'embedding', 1536, 'cosine')
            """,
            label="Message",
        )
    return

インデックスの確認

インデックスが作成できたかどうかは、以下のクエリで確認できます。

SHOW VECTOR INDEXES
YIELD name, type, labelsOrTypes, properties, options

インデックス機能自体は、以前からあったものなので、SHOW INDEXESでも表示できます。

SHOW INDEXES
YIELD name, type, labelsOrTypes, properties, options
WHERE type = 'VECTOR'

def show_index() -> list[str]:
    """NEO4jのインデックスを確認して、インデックス名をリストで返す"""
    with driver.session() as session:
        results = session.run(
            """
            SHOW INDEXES
            YIELD name, type, labelsOrTypes, properties, options
            WHERE type = 'VECTOR'
            """
            )
        response = []
        for result in results:
            response.append(result["name"])
        return response

インデックスの作成に成功していれば、以下のように作成したインデックスの情報が確認できます。（画像はNeo4j auraDBで直接クエリを打ち込んだものです）

2_ノードの作成

ベクトルデータの作成

ノードを作成するクエリを送る前に、ベクトルデータを作成する必要があります。
ここでは、OpenAIのtext-embedding-ada-002を使用しています。

from openai import OpenAI

client = OpenAI()

def get_embedding(text: str) -> list[float]:
    model = "text-embedding-ada-002"
    result = client.embeddings.create(input=[text], model=model).data[0].embedding
    return result

ベクトルを含むノードの作成

基本的なクエリは以下の通りです。

CREATE (m:Message {message: $message})
WITH m
CALL db.create.setNodeVectorProperty(m, 'embedding', $vector)

• CREATE (m:Message {name: $message})で新しいノードを作成し、

• WITH mでその変数を引き継いでいます。

• CALL db.create.setNodeVectorProperty()
  で、ベクトルの値をノードのプロパティに更新しています。
  各引数は以下の通りです。

  1. ノードを表す変数
  2. 'embedding'
  3. ベクトルの値（ここでは1536次元のlist[float]]）

コードでは、後で日付によるフィルタリングを行うため、ノードのプロパティに日時のデータも追加しています。

def create_message(message: str):
    """メッセージノードを作成する"""
    # メッセージテキストでベクトル作成
    vector = get_embedding(message)

    # 現在のUTC日時を取得し、ISO 8601形式の文字列に変換
    current_utc_datetime = datetime.utcnow()
    current_time = current_utc_datetime.isoformat() + "Z"

    with driver.session() as session:
        session.run(
            """
            CREATE (m:Message {name: $message, create_time: datetime($create_time)})
            WITH m
            CALL db.create.setNodeVectorProperty(m, 'embedding', $vector)
            """,
            message=message,
            create_time=current_time,
            vector=vector,
        )
    return

試しに、20個ほど適当な文章でノードを作りました。
作成されたノードのプロパティは、こんな感じに格納されます。

3_ベクトル検索の実行

作成したベクトルインデックスの検索は、以下のクエリで実行できます。

CALL db.index.vector.queryNodes('Message', $k, $vector)
YIELD m, score
RETURN m.name as message, score

各引数は以下の通りです。

1. インデックス名
2. 検索結果の個数
3. クエリのベクトルの値（ここでは1536次元のlist[float]]）

ここではノードの名前だけを取得していますが、返り値を、
RETURN properties(m) as properties
のようにすることで、ノードの他のプロパティも得ることができます。

ただし、その場合はベクトルの値（embedding）も含まれるので、うっかりそのままログに流してしまうと、ログが溢れます。

def query_messages(query: str, k: int = 3):
    """ベクター検索でメッセージノードを検索する。"""
    vector = get_embedding(query)

    with driver.session() as session:
        result = session.run(
            """
            CALL db.index.vector.queryNodes('Message', $k, $vector)
            YIELD node, score
            RETURN node.name as message, score
            """,
            k=k,
            vector=vector,
        )

        messages = []
        for record in result:
            message = {"message": record["message"], "score": record["score"]}
            if message:
                print(message)
                messages.append(message) if message else None
    return messages

ベクトル検索の結果

試しに、以下のようなクエリを実行したところ、ちゃんと意味の近い文章の結果が得られました。

query = "オフィスワーカーの仕事ってどんな感じですか？"
query_messages(query, k=3)

{'message': 'コンピューターと大きな暗い窓の前に２人が座っています。人々はオフィスワーカーです。', 'score': 0.9300100803375244}
{'message': 'ラップトップコンピューターを使用して机に座っている若い白人男。黒人はデスクトップコン ピューターを使用します。', 'score': 0.9017149806022644}
{'message': 'ひざの上に衣服を着て椅子に座っている年配の男性。高齢者が外に座っています。', 'score': 0.8990094065666199}

ベクトル検索の実行（追加条件を付ける）

さらに、ベクトル検索の結果に追加条件を付け、フィルタリングも行ってみました。

残念ながら、ベクトル検索自体の構文
CALL db.index.vector.queryNodes()にフィルターをつけることはできないのですが、ベクトル検索の結果にWHERE等の条件式を追加することで、特定の時期のメッセージだけを結果として抽出して返すようにすることができます。

例えば、以下のクエリでは、
ベクトル検索では件数を増やして広く結果を獲得し、365日（1年）以内の結果に限定し、スコア（ベクトル検索の類似度）が0.9以上のものを、スコアの高いものから順番にk個 抽出するようにしています。
(ORDER BYしなくても、ベクトル検索の結果はスコア順に出てきます。）

def query_messages_with_filter(query: str, k: int = 3):
    """ベクター検索でメッセージノードを検索し、フィルタリングする。"""
    vector = get_embedding(query)

    threshold: float = 0.9
    time_threshold: int = 365

    with driver.session() as session:
        result = session.run(
            """
            CALL db.index.vector.queryNodes('Message', $init_k, $vector)
            YIELD node, score
            WHERE score > $threshold AND node.create_time > datetime() - duration({days: $time_threshold})
            WITH node, score
            LIMIT $k
            RETURN node.name as message, score
            """,
            init_k=k * 10,
            k=k,
            vector=vector,
            threshold=threshold,
            time_threshold=time_threshold,
        )

        messages = []
        for record in result:
            message = {"message": record["message"], "score": record["score"]}
            if message:
                messages.append(message) if message else None
    return messages

4_まとめ

こんな感じで、Neo4jでは、グラフデータベースに対して、ベクトル検索を用いることができます。ぜひ皆さんも使ってみてください。