[Neptune Analytics x Llama-Index] ベクトル検索の改善

# 概要

Neptune AnalyticsとLlama-Indexライブラリを使用し、ベクトル検索を行おうとした際に、思うような挙動をしなかったため修正した概要を記載

# 各種リンク

[Neptune Analytics](https://docs.aws.amazon.com/ja_jp/neptune-analytics/latest/userguide/what-is-neptune-analytics.html)

[Llama-Index](https://www.llamaindex.ai/)

[Llama-Index(Github)](https://github.com/run-llama/llama_index)

※実行環境はSageMaker Notebookですが、環境構築などは別の記事をご参照ください。

# 目次

- 現状
- 課題
- 改善点
- 修正後

## 現状

### 目的
    
クエリキーワードに近い意味の言葉のノードを検索する

### 要点
- Llama-Indexのプロパティグラフを用いてグラフを構築
    - https://github.com/run-llama/llama_index/tree/main/llama-index-core/llama_index/core/indices/property_graph
- 検索時は Retrieverである`VectorContextRetriever`を使用
    - https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/vector.py
- `VectorContextRetriever` で検索を行うと、最終的に以下ファイルの`vector_query`関数でNeptune Analyticsに対してクエリを実行する
    - https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/graph_stores/llama-index-graph-stores-neptune/llama_index/graph_stores/neptune/analytics_property_graph.py#L73

下記コードは`vector_query`関数のクエリ実行箇所
```python
data = self.structured_query(
            f"""
            MATCH (e:`{BASE_ENTITY_LABEL}`)
            WHERE ({filters})
            CALL neptune.algo.vectors.get(e)
            YIELD embedding
            WHERE embedding IS NOT NULL
            CALL neptune.algo.vectors.topKByNode(e)
            YIELD node, score
            WITH e, score
            ORDER BY score DESC LIMIT $limit
            RETURN e.id AS name,
                [l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
                e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                score""",
            param_map={
                "embedding": query.query_embedding,
                "dimension": len(query.query_embedding),
                "limit": query.similarity_top_k,
            },
        )
```

## 課題

上記クエリの内容はざっと以下の挙動となる

#### 引数
- embedding：クエリキーワードをベクトル化したもの
- dimension：embeddingの次元数
- limit：ベクトル検索結果の取得件数（SQLのlimitのようなもの）

#### 挙動
1. ノードを取得し、各ノードのベクトルを取得
1. 取得した各ノードのベクトル値と近いベクトルを持つノードを検索し、ノードとスコアを取得
1. 全体のスコアを降順に整列
1. 結果の整形

### 課題点

- 引数の`embedding`が使用されていない
- `neptune.algo.vectors.topKByNode(e)`の使用が不適切
- スコアのソートが不適切

## 改善点

各課題の改善点は以下の通り

### 引数の`embedding`が使用されていない

本来クエリパラメータをクエリ内で使用するには、`$embedding`として使用してあげる必要があるが、当該クエリでは一切使用されている様子が見られない。

### `neptune.algo.vectors.topKByNode(e)`の使用が不適切

[公式ドキュメント](https://docs.aws.amazon.com/ja_jp/neptune-analytics/latest/userguide/vectors-topKByNode.html)では、下記引用のように記されている。（deeplによる翻訳）
> .vectors.topKByNodeアルゴリズムは、ノードからのベクトル埋め込み距離に基づいて、ノードの上位K個の最近傍を見つける。

つまり、当該クエリにおいては、ベクトルを持つ各ノードと近いベクトルを持つノードとそのスコアを取得していることになる。
この挙動は今回の目的からは逸脱してしまっているため、`.vectors.topKByEmbedding`を使用する。

[公式ドキュメント](https://docs.aws.amazon.com/ja_jp/neptune-analytics/latest/userguide/vectors-topKByEmbedding.html)では、下記のように述べられている。（deeplによる翻訳）
> .vectors.topKByEmbeddingアルゴリズムは、埋め込みベクトルの距離に基づいて、埋め込みベクトルの上位K個の最近傍を見つける。

つまり、このアルゴリズムの引数にクエリパラメータの`$embedding`を渡せば、上位K件のベクトルに近いノードを検索することが可能になる。

```cypher
WITH $embedding AS embedding
.vectors.topKByEmbedding(embedding)
```

### スコアのソートが不適切

[公式ドキュメント](https://docs.aws.amazon.com/ja_jp/neptune-analytics/latest/userguide/vectors-distance.html)では、Neptune Analyticsのベクトル検索のスコアについて以下の様に述べられている(deeplによる翻訳)
> .vectors.distanceアルゴリズムは，埋め込みに基づいて2つのノード間の距離を計算します．この距離は，埋め込みベクトルのL2ノルムの2乗です．

つまり、「スコアが小さいほど類似している」といえる。
そのため、当該クエリでは「昇順」が正しい。

## 修正後

上記改善点を踏まえると以下のクエリとなる。

:::note warn
注意
筆者はこのクエリに条件などを追加して動作確認しているため、このクエリ単体で動作するかは未検証です。
ご使用になる際は、別途動作確認をお願いします。
:::

```diff_python
data = self.structured_query(
            f"""
-           MATCH (e:`{BASE_ENTITY_LABEL}`)
-           WHERE ({filters})
-           CALL neptune.algo.vectors.get(e)
-           YIELD embedding
-           WHERE embedding IS NOT NULL
-           CALL neptune.algo.vectors.topKByNode(e)
-           YIELD node, score
-           WITH e, score
+           WITH $embedding AS embedded_query
+           CALL neptune.algo.vectors.topKByEmbedding(embedded_query)
+           YIELD embedding, node, score
+           WITH node, score
-           ORDER BY score DESC LIMIT $limit
-           RETURN e.id AS name,
-               [l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
+           ORDER BY score ASC LIMIT $limit
+           RETURN node.id AS name,
+               [l in labels(node) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
                e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                score""",
            param_map={
                "embedding": query.query_embedding,
-               "dimension": len(query.query_embedding),
                "limit": query.similarity_top_k,
            },
        )
```