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

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


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

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

https://neo4j.com/docs/getting-started/

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

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

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

https://neo4j.com/blog/vector-search-deeper-insights/

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

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

https://neo4j.com/cloud/platform/aura-graph-database/

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

① Neo4j auraに登録する。<br> 
② インスタンスを立ち上げる。（1個まで無料）<br>
③ pythonから使う場合は、neo4jのドライバを入れて、クエリを送信する。
（ブラウザから直接クエリを入力することもできる。）<br>
④ ドライバの初期化を行うコード<br>

```python
from neo4j import GraphDatabase

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

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

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

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


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

[Vector search index - Cypher Manual](https://neo4j.com/docs/cypher-manual/current/indexes-for-vector-search/)

## 1_ベクトルインデックスの作成
まずNeo4jのベクトルインデックスを作成します。
クエリは以下の通りです。
```
CALL db.index.vector.createNodeIndex(
    'Message', 'Message', 'embedding', 1536, 'cosine')
```
ここで、各引数は以下の通りです。
1. インデックス名　(何でも構いません）
1. ラベル名
1. 'embedding'
1. ベクトルの次元数（ここではopenAIのembeddingを使用）
1. ベクトルの比較手法

```python
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'
```
```python
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で直接クエリを打ち込んだものです）
![スクリーンショット 2024-01-13 191944.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2356928/f98449b7-a49d-f9a7-6ec1-6ca30d5dc961.png)


## 2_ノードの作成
#### ベクトルデータの作成
ノードを作成するクエリを送る前に、ベクトルデータを作成する必要があります。
ここでは、OpenAIの`text-embedding-ada-002`を使用しています。
```python
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. ノードを表す変数
    1. 'embedding'
    1. ベクトルの値（ここでは1536次元のlist[float]]）

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

```python
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個ほど適当な文章でノードを作りました。
作成されたノードのプロパティは、こんな感じに格納されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2356928/cbb7f308-1832-0067-aae6-9579be62348e.png)


## 3_ベクトル検索の実行
作成したベクトルインデックスの検索は、以下のクエリで実行できます。
```
CALL db.index.vector.queryNodes('Message', $k, $vector)
YIELD m, score
RETURN m.name as message, score
```

各引数は以下の通りです。
1. インデックス名
2. 検索結果の個数
2. クエリのベクトルの値（ここでは1536次元のlist[float]]）

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

ただし、その場合はベクトルの値（embedding）も含まれるので、うっかりそのままログに流してしまうと、ログが溢れます。
```python
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
```

#### ベクトル検索の結果
試しに、以下のようなクエリを実行したところ、ちゃんと意味の近い文章の結果が得られました。
```python
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しなくても、ベクトル検索の結果はスコア順に出てきます。）


```python
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では、グラフデータベースに対して、ベクトル検索を用いることができます。ぜひ皆さんも使ってみてください。