Supabase + Next.jsアプリ開発(2)-ベクトルDB

2024.07.23

各社の自社機密情報を参照するQ&Aチャットボットのようなアプリを想定します。ChatGPTのAPIを使用するとして、ChatGPTの学習データにはない各社機密情報についての回答をさせるには、RAG(Retrieval-Augmented Generation)という技術が使われます。詳細は割愛しますが、この技術を使うには「機密情報」を「ベクトル化(Embedding)」しておいて、ユーザーの要請に関連する情報をベクトル検索してChatGPTに提示してあげることが必要で、このアプリで使用するDBは「ベクトル」を扱えることが必須要件となります。

幸いなことに、PostgreSQLにはpgvectorという拡張機能があり、Supabaseでも利用できるようになっています。このTopicではその辺りの設定について公式サイトを元に解説していきます。

pgvector: Embeddings and vector similarity | Supabase Docs
pgvector: a PostgreSQL extension for storing embeddings and performing vector similarity search.
supabase.com

前回Topicで作成したSupabaseダッシュボードにて

  • 左メニューで「Database」を選択
  • 左メニューで「Extensions」を選択
  • 「Available extensions」から「vector」を探し、有効化

以上の準備だけで「pgvector」が使用できるようになります。OpenAIのEmbeddingモデルの１つである「text-embedding-3-small」を用いるとベクトルの次元は1536になるので、以下のクエリでテーブルを作成します。

create table embeddings (
  id serial primary key,
  title text not null,
  body text not null,
  embedding vector(1536)
);

作成したテーブルに、次のようなデータを入力しておきます。（ちなみにすでにembeddingのところに値が入ってしまっていますが、ここはNullで良いです。この後OpenAIのAPIを使ってembeddingを計算してテーブルをupdateします）。

次に、OpenAIのAPIを使って情報をベクトル化してテーブルに保存するコード例を示します。前回のTopicでNext.jsのアプリは作成ずみとします。

import { supabase } from "@/utils/supabase/supabase"
import OpenAI from "openai"

const openai = new OpenAI({    apiKey: process.env.OPENAI_API_KEY,  });

let { data: stories, error } = await supabase
          .from('embeddings')
          .select()
        if (stories?.length == 0){
            console.log("no data")
        } else {
            for (const story of stories!){
                if (!story.embedding){
                    const embedd = await openai.embeddings.create({
                        model: "text-embedding-3-small",
                        input: story.body,
                        encoding_format: "float",
                    });
                    const vector = embedd.data[0].embedding
                    const { error } = await supabase
                    .from('embeddings')
                    .update({embedding:vector})
                    .eq('id',story.id)
                }
            }
        }

これで、bodyの内容に相当するベクトルが生成され、データベースを更新することができました。