Supabase+pgvector+OpenAI API+Remixで類似度検索を実装してみる
はじめに
こんにちは、健常者エミュレータ事例集の管理人をしているcontradiction29です。
「健常者エミュレータ事例集」は、「個人の属性に寄らず、誰もが暗黙知を投稿でき、閲覧でき、評価できるプラットフォームを作る」をコンセプトに開発が進められているプロジェクトです。以下のリンクからアクセスできるので、よかったら閲覧してみてください。
ユーザーはテンプレートにそって経験を整理し、知識として共有し、自由に評価し、コメントで議論ができます。GPL3.0ライセンスの元、コード自体も公開されています。詳しくは以下のレポジトリをご覧ください。
健常者エミュレータ事例集には、今読んでいる記事と関連した記事を表示する機能が実装されています。例えば、記事食事中に料理の味を悪く言うのはよくないの関連記事は以下のようになっています。
※ 2024-04-12時点です
今回は、この関連記事の推薦機能の裏側を紹介しつつ、類似度検索を実装したシンプルな投稿フォームを作ってみようと思います。使う技術は以下の通りです。
-
OpenAI embedding API
- 今回の目玉です。記事の中身を入力し、embeddingした結果を返します。
- ほかのAPIでも代用は可能ですが、ローカルLLMを除いたらこれが一番安いと思います。
- ちなみに「健常者エミュレータ事例集」に投稿された記事の平均トークン数は633程度であり、最大トークン数が8192もあるtext-embedding-3-smallだと3つを除いたほぼすべての記事をembeddingすることが可能でした(3つだけトークン数が過大になってダメだった）
-
Supabase
- embeddingした結果を格納するための場所、およびそれを取り出すための機構として使います。
- Supabase自体は様々なバックエンド機能の集合体ですが、その中にマネージドPostgresSQLとしての機能があります。拡張機能としてpgvectorを利用することができるため、ベクターストアとしての利用が可能となります。
- イメージとしては「PostgresSQLで作成したテーブルを構成するカラムの中にvector型を利用できるようになり、それをベクターストアとする」感じです。
- ほかのベクターストアでも代用は可能です。
-
Remix
- Reactをベースにしたフルスタックフレームワークです。
- 今回説明する仕組みは特段フレームワークに依存するものではないため、ほかのフレームワークでも代用は可能です。
それでは中身に入っていきます。
参考記事
OpenAI Embedding APIに関しては、公式ドキュメントが一番よくまとまっています。
Supabaseを使って類似度検索を実装した際に参考にさせていただきました。
Remixを学習する際に利用しました。オライリーのサブスクにも入っています。
ざっくりと仕組みを解説
実際に健常者エミュレータ事例集で行われている類似度検索の全体の流れを図でざっくりと示します。
流れは以下の通りです：
- フォーム上にユーザーが入力したデータをサーバーサイドで受け取る
- サーバーサイド側で受け取ったデータをOpenAI Embedding APIに投げ、結果をSupabaseに格納する
- PostgresSQL(Supabase)のfunction機能を利用し、類似する記事を取得して返す
それでは、この流れを一緒に実装してみましょう。
細かい仕組み
実装を細かく見ていきます。Node.jsランタイムを想定しています。
前提
- Remixに関する知識はあるものとします
- Supabaseのアカウント作成は済んでいるものとします
- OpenAIのAPI Keyの取得は済んでいるものとします
-
.envにOPENAI_API_KEY, SUPABASE_URL, SUPABASE_ANON_KEYが格納されているとします
npx create-remix@latest // remixのテンプレートで初期化します
cd ${app_name} // 作成したディレクトリに移動
フォームで入力したデータを受け取る
/app/routes配下に以下のファイルを作成します。シンプルな投稿フォームです。
import { ActionFunctionArgs } from "@remix-run/node";
import { Form } from "@remix-run/react";
export async function action({ request }:ActionFunctionArgs) {
const formData = await request.formData();
const text = formData.get("text");
console.log(text);
return new Response("Thanks for submitting!");
}
export default function SampleForm(){
return (
<>
<h1>Sample Form</h1>
<Form method="post">
<label>
感情を入力してください:
<input type="textarea" name="text" />
</label>
<button type="submit">Submit</button>
</Form>
</>
)
}
実際にフォームを動作させ、コンソールに入力した文字が出力されることを確認してみましょう。
npm run dev
フォームのデータをOpenAI Embedding APIに投げ、結果をSupabaseに格納する
/app配下にmodulsディレクトリを作成し、/app/modules配下に以下のファイルを作成します。
import { json } from "@remix-run/node";
const OpenAIAPIKey = process.env.OPENAI_API_KEY;
const OpenAIEmbeddingEndpoint = "https://api.openai.com/v1/embeddings"
const OpenAIEmbeddingModel = "text-embedding-3-small"
export async function createEmbedding(text:string){
if (!OpenAIAPIKey){
throw new Error("OpenAI API Key not set");
}
const response = await fetch(OpenAIEmbeddingEndpoint, {
method: "POST",
headers: {
"Content-Type": "application/json",
"Authorization": `Bearer ${OpenAIAPIKey}`
},
body: JSON.stringify({
model: OpenAIEmbeddingModel,
input: text,
encoding_format : "float"
})
});
if (!response.ok){
throw new Error("Failed to create embedding");
}
const parsedResponse = await response.json();
const embedding = Array.from(parsedResponse.data[0].embedding);
const tokenCount = parsedResponse.usage.total_tokens
console.log(embedding, tokenCount)
return json({embedding, tokenCount});
}
処理の流れはシンプルで、OpenAIのembedding API用のエンドポイントにデータを投げ、その結果を取得しているだけの処理になっています。
embedding.server.tsを先ほどのフォームに組み込みます。
import { ActionFunctionArgs } from "@remix-run/node";
import { Form } from "@remix-run/react";
import { createEmbedding } from "~/modules/embedding.server";
export async function action({ request }:ActionFunctionArgs) {
const formData = await request.formData();
const text = formData.get("text");
console.log(text);
if (!text){
return new Response("Please enter some text", { status: 400 });
}
const embeddingResult = await createEmbedding(text.toString());
return new Response("Thanks for submitting!");
}
export default function SampleForm(){
return (
<>
<h1>Sample Form</h1>
<Form method="post">
<label>
感情を入力してください:
<input type="textarea" name="text" />
</label>
<button type="submit">Submit</button>
</Form>
</>
)
}
フォームからテキストを投稿し、コンソールにベクトルとトークン数が出力されれば成功です。
なお、トークン数は本来必要ない情報ですが、実際にシステムに組み込む際にコストの見積もりに使うのでここで取得しています。
次に、このデータをSupabaseに格納するため、embedding.server.tsを書き換えます。
SupabaseのSQL Editorで以下のようなテーブルを作成してください。拡張機能を有効化し、anonロールに対する権限も付与しておきます。
create extension vector with schema extensions; -- 拡張機能pgvectorを有効化
drop table if exists sample_table;
create table
sample_table (
id int primary key generated always as identity,
content text,
embedding vector (1536),
token_count int
);
grant usage on schema public to anon;
grant insert on table sample_table to anon;
embedding.server.tsに、SupabaseにデータをInsertする処理を追加します。
import { json } from "@remix-run/node";
import { createClient } from "@supabase/supabase-js";
const OpenAIAPIKey = process.env.OPENAI_API_KEY;
const OpenAIEmbeddingEndpoint = "https://api.openai.com/v1/embeddings"
const OpenAIEmbeddingModel = "text-embedding-3-small"
const SUPABASE_ANON_KEY = process.env.SUPABASE_ANON_KEY;
const SUPABASE_URL = process.env.SUPABASE_URL;
export async function createEmbedding(text:string){
if (!OpenAIAPIKey){
throw new Error("OpenAI API Key not set");
}
const response = await fetch(OpenAIEmbeddingEndpoint, {
method: "POST",
headers: {
"Content-Type": "application/json",
"Authorization": `Bearer ${OpenAIAPIKey}`
},
body: JSON.stringify({
model: OpenAIEmbeddingModel,
input: text,
encoding_format : "float"
})
});
if (!response.ok){
throw new Error("Failed to create embedding");
}
const parsedResponse = await response.json();
const embedding = Array.from(parsedResponse.data[0].embedding);
const tokenCount = parsedResponse.usage.total_tokens
if (!SUPABASE_ANON_KEY || !SUPABASE_URL){
throw new Error("Supabase credentials not set");
}
const supabaseClient = createClient(SUPABASE_URL, SUPABASE_ANON_KEY);
const { data, error } = await supabaseClient.from("sample_table").insert([
{ content: text, embedding: embedding, token_count: tokenCount }
]);
if (error){
console.log(error)
throw new Error("Failed to insert into database");
}
console.log(data)
return json({embedding, tokenCount});
}
こちらの処理もシンプルで、受け取ったデータをSupabase上のテーブルにinsertしているだけの処理になります。
これで、embeddingした結果とトークンカウント数、テキスト内容がSupabase内部で保持されるようになりました。権限周りは躓きやすいので、エラーが起こっている場合は権限が適切に付与されているか確認してみてください。
PostgresSQL(Supabase)のfunction機能を利用し、類似する記事を取得して返す
以下のように関数を作成します
create or replace function sample_search_similar_content (
query_text text,
match_count int
)
returns table(
id int,
content text,
similarity float
)
language sql stable
as $$
select
p.id,
p.content,
(p.embedding <#> q.embedding) * -1 as similarity
from sample_table p
cross join (
select embedding
from sample_table
where content = query_text
) q
order by similarity desc
limit match_count;
$$;
grant execute on function sample_search_similar_content to anon;
grant select on table sample_table to anon;
この関数をaction関数内部で実行し、Submitボタンを押したら過去の投稿で似たようなものが表示されるようにしてみます。supabaseで登録されたfunctionを利用するためには、supabaseClient.rpcメソッドを利用します。引数としてfunctionで設定したパラメータを受け渡します。
import { ActionFunctionArgs, json } from "@remix-run/node";
import { Form, useActionData } from "@remix-run/react";
import { createClient } from "@supabase/supabase-js";
import { createEmbedding } from "~/modules/embedding.server";
const SUPABASE_ANON_KEY = process.env.SUPABASE_ANON_KEY;
const SUPABASE_URL = process.env.SUPABASE_URL;
export async function action({ request }:ActionFunctionArgs) {
const formData = await request.formData();
const text = formData.get("text");
console.log(text);
if (!text){
return new Response("Please enter some text", { status: 400 });
}
await createEmbedding(text.toString());
if (!SUPABASE_ANON_KEY || !SUPABASE_URL){
throw new Error("Supabase credentials not set");
}
const supabaseClient = createClient(SUPABASE_URL, SUPABASE_ANON_KEY);
const { data, error } = await supabaseClient.rpc("sample_search_similar_content", { query_text: text, match_count: 11 });
if (error){
console.log(error)
throw new Error("Failed to search similar content");
}
const response = data.slice(1,11)
return json({ response });
}
export default function SampleForm(){
const actionData = useActionData<typeof action>();
return (
<>
<h1>Sample Form</h1>
<Form method="post">
<label>
感情を入力してください:
<input type="textarea" name="text" />
</label>
<button type="submit">Submit</button>
</Form>
<h2>Similar Content</h2>
<ul>
{actionData && actionData.response.map((content:any) => (
<li key={content.id}>{content.content}</li>
))}
</ul>
</>
)
}
先ほどの関数を用いて類似度検索を行う場合、最も類似度が高いのはその記事自身になるため、それを除外しています。
ところでなんですが、今日（2024年4月12日)は散歩をしていたら桜がきれいでした。満開の時期は過ぎましたが、葉桜も美しいですね。
そういうわけなので、百人一首のうち最初の20句をサンプルデータとして入れてみます。
参考：
Hey, OpenAI 「散ればこそいとど桜はめでたけれうき世になにか久しかるべき」に近い句は？
まぁまぁやるじゃん
終わりに
実際はもう少し複雑な実装をしています。詳しくは以下のコードをご覧ください。
また、インデックスを張らないと実用に耐えうる計算速度にはならないので、以下のようなSQLでインデックスを作成することを推奨します。
CREATE INDEX ON sample_table USING hnsw (embedding vector_ip_ops);
健常者エミュレータ事例集では開発を手助けしてくれる方を募集しています。よかったら覗いてみてください。
Discussion