Firestoreのベクトル検索で全文検索してみる

# 目的

Firestoreの全文検索の実装にベクトル検索を使ってみて良かったので良さを共有したい

## 背景
趣味の個人開発でデータベースサイトを作りたいと思い立ち、そのためのデータベース基盤としてFirebaseのFirestoreを使うことにしました。

そこでテキストの全文検索機能を設計するにあたり、Firestoreの標準のクエリだと以下のような理由で厳しかったので何かしら別の策を考える必要が生じました。

- Firestoreのクエリはドキュメント単位で課金されるため、ドキュメント総当たりでフィルタリングするとすれば課金額が膨大になる
- Firestoreのクエリ演算子で検索に使えるものには完全一致（`==`）、配列内一致（`array-contains`）しかなく、検索キーワードとの類似度でソートするといったことができない
- Elastic Search、Algoliaといった外部サービスを使えば一応全文検索できるが、これも結構コストがそこそこ掛かる

　
色々調べているうちに、Firestoreに最近追加された**ベクトル検索**を使えば行けそうだと分かったので、これを試してみることにしました。

https://zenn.dev/miyasic/articles/firebase-search-with-vector-embeddings

https://firebase.google.com/docs/firestore/vector-search?hl=ja

# ベクトル検索について

### **ベクトル検索**とは

> ベクトル検索とは、データを高次元のベクトル空間にマッピングし、類似性に基づいて情報を検索する技術です。
[ZOZO TECH BLOGより](https://techblog.zozo.com/entry/vector-search-introduction#%E3%83%99%E3%82%AF%E3%83%88%E3%83%AB%E6%A4%9C%E7%B4%A2%E3%81%A8%E3%81%AF)
> 
ざっくり噛み砕くと、テキストや画像などのデータを数字の集合（[-0.00949603, 0.021038500, -0.01994903, …]のような多次元行列）に変換し、集合同士の近似度を内積などで計算し、より近いデータを見つける手法です。

検索以外にもレコメンドシステムでも使われているようなので、応用性の高い技術だと言えます。

https://engineering.mercari.com/blog/entry/20230612-cf-similar-item/

[協調フィルタリングとベクトル検索エンジンを利用した商品推薦精度改善の試み | メルカリエンジニアリング](https://engineering.mercari.com/blog/entry/20230612-cf-similar-item/)
    

### **K最近傍（KNN）法**

- ベクトル検索のアルゴリズムの一手法です。詳しくは↓
    

    :::note info
    **ChatGPTによる解説**
    
    K近傍法（K-Nearest Neighbors, KNN）は、分類や回帰に使われるシンプルな機械学習アルゴリズムです。あるデータを分類する際、そのデータに最も近いK個の既知のデータ（近傍）のクラスを参照して、多数決で新しいデータのクラスを決定します。例えば、K=3の場合、最も近い3つのデータ点のクラスのうち多数を占めるクラスに新しいデータを分類します。
    
    KNNの特徴は、距離（例えばユークリッド距離など）を基にして近いデータを見つけるため、訓練時にモデルを作成するのではなく、予測時に計算を行う「遅延学習」の手法であることです。また、Kの値の選び方やデータのスケーリングが性能に影響を与えます。

### ベクトル化について

- テキストを意味や類似度で分類されたベクトルに変換するには機械学習モデルが必要です
    ```mermaid
    graph LR
    A[ユーザーからのテキスト入力] --> B[機械学習モデルでの処理]
    B --> C[ベクトル化されたテキスト]
    ```
- クラウド上でPythonを動かせるFirebase Cloud Functionsに機械学習用のフレームワーク（`torch`、`transformers`など）とモデルをデプロイしようとしたところ、Cloud Functionsのメモリ制限に引っ掛かってしまい断念
- 下記ブログでOpenAI APIのベクトル埋め込み用の`text-embedding-3-small`モデルを使う方法が紹介されており、Functionsにフレームワークやモデルをデプロイしなくて良い（容量が節約できる）のと、コスト面で非常に安価[^2]だったので今回はこれを採用しました

[^1]: 2024年10月25日現在$0.020/1Mトークン。100万文字のテキストを変換して約3円。

https://note.com/npaka/n/nfedf04b025ae

# いざ実装

### 全体フロー図

```mermaid
sequenceDiagram
    participant User as User
    participant CloudFunction as Cloud Function
    participant OpenAI as OpenAI API
    participant Firestore as Firestore

    User->>CloudFunction: 新しいテキストデータを送信
    CloudFunction->>OpenAI: テキストをベクトル化リクエスト
    OpenAI-->>CloudFunction: ベクトル化されたデータを返却
    CloudFunction->>Firestore: ベクトル化されたデータを保存

    User->>CloudFunction: 検索クエリ送信
    CloudFunction->>OpenAI: クエリをベクトル化
    OpenAI-->>CloudFunction: クエリベクトルを返却
    CloudFunction->>Firestore: ベクトル検索で類似データを取得
    Firestore-->>CloudFunction: 類似データの返却
    CloudFunction-->>User: 類似データの結果を返す

```
## Firebase

### 1. ReactのプロジェクトにFirebaseとFirestoreを追加

公式ドキュメントなどを参照してセットアップしてください。
こちらの記事でも分かりやすくまとめられています。

https://qiita.com/nemutas/items/330cab413bccd9b2e096

### 2. Cloud Functionsをセットアップ

下記記事などを参考に、各環境に合わせてセットアップしてください。

https://zenn.dev/yusu29/articles/firebase_cloudfunctions_split

Cloud Functionsは無料プランでは利用不可で、従量課金プラン（Blaze）に変更する必要があります。プランをアップグレードしても無料枠を越えなければお金はかかりませんが、コストを意識して運用するのも勉強になるので勇気を出してアップグレードしましょう。

### 3. Firestoreにベクトルインデックスを追加
Firestoreのコレクションに対してKNNベクトル検索のクエリを実行するにはコレクションのベクトルインデックスが作成されている必要があります。

任意のコレクションを作成して、以下の手順でベクトルインデックスを作成します。

#### ベクトルインデックス作成手順
- （前提）`google-cloud-sdk`をインストールしていること

https://qiita.com/G-awa/items/e6904b040caa0096fba0

- `gcloud firestore indexes composite create`コマンドでベクトルインデックスを作成（2,3分掛かる）
        
```bash
(base) user@MacBook-Air react-project % gcloud alpha firestore indexes composite create \
    --collection-group=インデックス対象コレクション名 \
    --query-scope=COLLECTION \
    --field-config field-path=embedding_field,vector-config='{"dimension":"1536", "flat": "{}"}' \
    --database="(default)"
```
        
- `--collection-group`：インデックスを新規作成するFirestoreのコレクション名
- `--field-config`
    - `field-path`：インデックス化したいベクトルデータが入ったフィールド名（今回は`embedding_field`とする）
    - `vector-config`
        - `dimension`：ベクトルの次元数。今回はOpenAIの`text-embedding-3-small`モデルを利用するので1536を指定（上限は2048）
        - `flat`：インデックスのタイプがflatであることを明記。他のオプションが無いのである意味おまじない。
    - `database`：対象とするデータベース。アプリにFirebaseプロジェクトが紐付けられていたら`"(default)"`でok
　
- 他のフィールドと掛け合わせた複合インデックスを作りたい場合は次のようなコマンドにする
```bash
gcloud alpha firestore indexes composite create \
    --collection-group=collection-group \
    --query-scope=COLLECTION \
    --field-config=order=ASCENDING,field-path="color" \
    --field-config field-path=vector-field,vector-config='{"dimension":"1024", "flat": "{}"}' \
    --database=database-id
```
- 作成したインデックスは`gcloud alpha firestore indexes composite list --database="(default)”`で確認できる
```bash
(base) user@MacBook-Air react-project % gcloud alpha firestore indexes composite list --database="(default)"
---
fields:
- fieldPath: __name__
  order: ASCENDING
- fieldPath: embedding_field
  vectorConfig:
    dimension: 1536
    flat: {}
name: projects/react-study/databases/(default)/collectionGroups/sampleCollection/indexes/CICAgOi3kJAK
queryScope: COLLECTION
state: READY
```

https://firebase.google.com/docs/firestore/vector-search?hl=ja#create_a_vector_index

## OpenAI API

### OpenAI APIキーを取得

こちらの記事を参考にAPIキーを取得してください。

https://qiita.com/kurata04/items/a10bdc44cc0d1e62dad3

Firebaseの従量課金プランと違い、こちらはプリペイドで料金をチャージし、使った分だけ残高から差し引かれる仕組みになってます。
大量にAPIリクエストをしない限り1円未満の請求で済む上、上限に達したところでAPIリクエストを受け付けなくなるだけです。[うっかり事故](https://gigazine.net/news/20180803-spent-much-money-in-firebase/)になる心配が無いのは魅力的ですね。

APIリクエストとレスポンスは以下のような形式です。

```bash:request
curl https://api.openai.com/v1/embeddings \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $OPENAI_API_KEY" \
  -d '{
    "input": "Your text string goes here",
    "model": "text-embedding-3-small"
  }'
```

```json:response
{
  "object": "list",
  "data": [
    {
      "object": "embedding",
      "index": 0,
      "embedding": [
        -0.006929283495992422,
        -0.005336422007530928,
        ...
        -4.547132266452536e-05,
        -0.024047505110502243
      ],
    }
  ],
  "model": "text-embedding-3-small",
  "usage": {
    "prompt_tokens": 5,
    "total_tokens": 5
  }
}
```

https://platform.openai.com/docs/guides/embeddings/how-to-get-embeddings

## Cloud Functionsでのベクトル検索
### ベクトル化したテキストをFirestoreに保存

- Cloud Functionsで以下の処理を含む関数を定義します
    1. OpenAI APIの`text-embedding-3-small`モデルを使い、テキストをベクトル化
    2. Firestoreのコレクションにベクトル化したテキストを保存
```jsx
const OPENAI_API_URL = "https://api.openai.com/v1/embeddings";
const OPENAI_API_KEY = "your_key_here";

// OpenAI APIコールとベクトル取得の共通関数
const getVectorEmbeddedText = async (text: string) => {
  const response = await axios.post(
    OPENAI_API_URL,
    { model: "text-embedding-3-small", input: text },
    {
      headers: {
        "Authorization": `Bearer ${OPENAI_API_KEY}`,
        "Content-Type": "application/json",
      },
    }
  );
  return response.data.data[0].embedding;
};

/* テキストをベクトル化するCloud Function関数（CORS設定は適当）*/
export const vectorEmbedText = onRequest(
  { cors: true }, 
  async (req, res) => {
  
      // POSTリクエストのボディからテキストを取得
      const { text } = req.body;
    
      try {
        // ベクトル化されたデータを取得
        const embedding = await getVectorEmbeddedText(text);
    
        // embeddingからFieldValue.vectorを生成
        const vector = FieldValue.vector(embedding);
    
        // コレクション 'sampleCollection' にドキュメントを追加
        const docRef = await db.collection("SampleCollection").add({
          text_field: text,
          embedding_field: vector,
          createdAt: admin.firestore.FieldValue.serverTimestamp(),
        });
        res.status(200).json({ success: true, id: docRef.id });
      } catch (error) {
        res.status(500).send("Error embedding text.");
      }
    }
);
```

上の関数を実行すると、Firestoreのコレクションにテキストとベクトルデータが追加されます。
成功すれば、下の画像のように`embedding_field`にベクトルデータが登録されていることを確認できます。

![スクリーンショット 2024-10-25 23.29.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/245143/32ada068-0c66-5c64-61ee-9dd2e1cc01d0.png)


### 検索キーワードをベクトル化し、最近傍クエリを実行して結果を取得
- こちらもCloud Functionsに次の関数を追加します
    1. OpenAI APIでテキストをベクトル化
    2. Firestoreのコレクションに対してベクトル検索クエリを実行
    
```jsx
/* 最近傍クエリを実行する関数（CORS設定は適当）*/
export const getNearestNeighbor = onRequest({ cors: true }, async (req, res) => {
    try {
      // GETリクエストのクエリからテキストを取得
      const text = req.query.text;
      if (typeof text !== "string") {
        throw new Error("Invalid request: 'text' query parameter is required.");
      }

      // ベクトル化されたデータを取得
      const embedding = await getVectorEmbeddedText(text);

      // embeddingからFieldValue.vectorを生成
      const vector = FieldValue.vector(embedding);
      
      // ベクトル検索クエリを実行
      const vectorQuery: VectorQuery = db.collection("sampleCollection")
                                         .findNearest({
                                           limit: 5,
                                           distanceMeasure: "EUCLIDEAN",
                                           vectorField: "embedding_field",
                                           queryVector: vector,
                                       });
      const snapshot: VectorQuerySnapshot = await vectorQuery.get();

      // レスポンスをお好みに処理
      ...
```

## 実際の動作

書籍タイトルを何個か登録して実際に検索クエリの結果を見てみました。

![vectorSearch.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/245143/9b254a19-0a8e-dc8b-8a6b-859210efb47a.gif)
　
登録データが少ないのでキーワードと無関係なドキュメントもヒットしますが、基本的には狙ったデータが上位に出るので非常に理想的な感じになりました🎉

Firestoreのドキュメント読み込み数も`limit`で絞った分だけカウントされているのでそこまで気にしなくていい数値に収まっていました。


# まとめ

### メリット
- 曖昧なキーワードでも期待通りのデータがヒットするのでとても良い✨
- Firestoreの標準のクエリと同じようにシンプルに書けるので開発者経験的にもGood

### デメリット
- コレクション全体の中から最も近いものをフェッチするので、期待するデータがないと無関係な検索結果ばかり出てくる
    - 最近、Amazonとかでキーワードと関係無い商品がヒットするのもベクトル検索使ってるから？🤔
- Firestoreの標準のクエリに備わった`offset`、`startAt`などクエリカーソルを操作するメソッドが使えないので、ページネーションを実現するには`limit`を増やして重複するデータを随時捌くしかない
    - （検索精度が高ければそもそもページネーション要らんのではと思わなくもない）

### その他の学び

- OpenAI APIの`text-embedding-3-small`モデルはお財布に優しい
    - 1000回検索して0.3円程度。コーヒー1杯で100万回検索できる☕

## 今後の課題
ベクトル検索は検索だけでなくレコメンドにも応用できるということだったので、今度はレコメンドエンジンの設計にも手を出してみたいです。

### 参考

https://firebase.google.com/docs/firestore/vector-search?hl=ja

https://note.com/npaka/n/nfedf04b025ae