LangChain JSを参考にしつつ、Supabaseで作成したベクトルデータベースの検索を実装してみる
はじめに
今回はタイトルの通り、Supabaseでベクトルデータベースを作成し、簡単なサンプルコードを動かす作業をします。その際、LangChainの実装を参考にしながら作ってみようと思います。
記事が長くなってしまったので、ベクトルデータベースやSupabaseの説明、これらを使う理由などについては以下の記事に書きましたので、必要であれば併せてご参照ください。
LangChain JSを参考にしつつ、Supabaseで作成したベクトルデータベースの検索を実装してみる（補足編）
前提
Supabaseのアカウント作成やAPIキーの取得については省きます。
開発言語はNode.jsを使い、ローカルのMac上でコードを実行しています。
LangChainのコード調査は2023.07.22にcloneしたものを用いています。
コミットID: df0dac452ac8c2ed6288a9285ee916e750f39b15
Supabase管理画面上からテーブルと関数の作成
まず、LangChainのドキュメント、Create a table and search function in your databaseに記載されている内容（テーブル作成、PostgresSQL関数の作成）を、Supabaseの管理画面から実行してしまいます。
pgvector拡張を読み込み
Supabaseの管理画面から以下を実行します。
-- Enable the pgvector extension to work with embedding vectors
create extension vector;
データベース内に新しい拡張、pgvectorを有効にします。
テーブル作成
-- Create a table to store your documents
create table documents (
id bigserial primary key,
content text, -- corresponds to Document.pageContent
metadata jsonb, -- corresponds to Document.metadata
embedding vector(1536) -- 1536 works for OpenAI embeddings, change if needed
);
contextには本文、embedding にはembeddingした本文を格納します。
metadataは、資料自体の名称や、例えば長文を分割して保存したような場合は、全体の中で位置（例えばページ数）など保存します。ここで設定した項目は通常のRDSの検索と同様に絞り込みに使えます。
metadata例：
{"loc":{"lines":{"to":357,"from":357}},"name":"坊っちゃん","author":"夏目漱石"}
例えば、"author"（著者名）、"name"（作品名）、"lines"（ページ番号)で絞り込みを加えることができます。
今回はLangChainのサンプル通り作成しますが、作成するアプリケーションによって独自のカラムを作っても良いかもしれませんね。
vector(1536)の1536はコメントにあるように次元数です。OpenAI API以外の方法でembeddingする場合は、ここを調整する必要があります。
注意点として、SupabaseのAPIキーは権限を絞ってフロントエンドで配置するような使い方をすることがあります。そういった場合、権限設定にはご注意ください。今回はローカル環境やサーバーサイドのような環境から使うので特別な設定を加えていません。
関数作成
-- Create a function to search for documents
create function match_documents (
query_embedding vector(1536),
match_count int DEFAULT null,
filter jsonb DEFAULT '{}'
) returns table (
id bigint,
content text,
metadata jsonb,
similarity float
)
language plpgsql
as $$
#variable_conflict use_column
begin
return query
select
id,
content,
metadata,
1 - (documents.embedding <=> query_embedding) as similarity
from documents
where metadata @> filter
order by documents.embedding <=> query_embedding
limit match_count;
end;
$$;
LangChainではSupabaseのクライアントから、このmatch_documentsという関数を呼び出して使います。
呼び出す際の引数は以下になります。
- query_embedding: 質問文をembeddingしたもの
- match_count: 取得件数
- filter: メタデータでの絞り込み
この実装では質問文（をembeddingしたもの）とcontentの類似度を計算し、似ているものから指定件数返却します。filterが指定されていればmetadataでも絞り込みを行います。
今回はこれを踏襲してmatch_documentsを呼び出す実装をしますが、今後作成するアプリケーションでは、呼び出し元で関数部分の実装を行うことも考えています。（呼び出し元でSQLを組み立てる）
LangChainでの実装例
ここから実装に移りますが、まずはLangChainでのサンプルコードを確認します。
import { SupabaseVectorStore } from "langchain/vectorstores/supabase";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { createClient } from "@supabase/supabase-js";
// First, follow set-up instructions at
// https://js.langchain.com/docs/modules/indexes/vector_stores/integrations/supabase
const privateKey = process.env.SUPABASE_PRIVATE_KEY;
if (!privateKey) throw new Error(`Expected env var SUPABASE_PRIVATE_KEY`);
const url = process.env.SUPABASE_URL;
if (!url) throw new Error(`Expected env var SUPABASE_URL`);
export const run = async () => {
const client = createClient(url, privateKey);
const vectorStore = await SupabaseVectorStore.fromTexts(
["Hello world", "Bye bye", "What's this?"],
[{ id: 2 }, { id: 1 }, { id: 3 }],
new OpenAIEmbeddings(),
{
client,
tableName: "documents",
queryName: "match_documents",
}
);
const resultOne = await vectorStore.similaritySearch("Hello world", 1);
console.log(resultOne);
};
処理内容としては、
まず、テキストとそれに紐付いたメタデータ、Embeddingに使うオブジェクト、Supabaseのクライアントと設定をSupabaseVectorStore.fromTexts() に渡し、これが
テキストをembeddingし、メタデータと一緒にSupabaseのベクトルデータベースに投入しています。
その後、vectorStore.similaritySearch()で検索テキストをembeddingしてから、ベクトルデータベースに類似度検索をかけ、最も類似度が高いレコード１件を取得しています。
データの投入部分の処理を読む
個別の処理を詳しく見ていきます。
データの投入処理は以下になります。
const vectorStore = await SupabaseVectorStore.fromTexts(
["Hello world", "Bye bye", "What's this?"],
[{ id: 2 }, { id: 1 }, { id: 3 }],
new OpenAIEmbeddings(),
{
client,
tableName: "documents",
queryName: "match_documents",
}
);
これだけだと何をしているのかわからないので、"SupabaseVectorStore.fromTexts()"の実装を見てみます。
あまりじっくりとコードを読まず、イメージを掴んでいただく程度で良いです。
static async fromTexts(
texts: string[],
metadatas: object[] | object,
embeddings: Embeddings,
dbConfig: SupabaseLibArgs
): Promise<SupabaseVectorStore> {
const docs = [];
for (let i = 0; i < texts.length; i += 1) {
const metadata = Array.isArray(metadatas) ? metadatas[i] : metadatas;
const newDoc = new Document({
pageContent: texts[i],
metadata,
});
docs.push(newDoc);
}
return SupabaseVectorStore.fromDocuments(docs, embeddings, dbConfig);
}
引数を加工して、"SupabaseVectorStore.fromDocuments()"を呼び出しています。
"Document"はLangChainでよく使われる、テキスト（コンテンツ）とメタデータをまとめたデータ構造です。
static async fromDocuments(
docs: Document[],
embeddings: Embeddings,
dbConfig: SupabaseLibArgs
): Promise<SupabaseVectorStore> {
const instance = new this(embeddings, dbConfig);
await instance.addDocuments(docs);
return instance;
}
続けて引数を使って "addDocuments()"が呼び出されます
async addDocuments(documents: Document[], options?: { ids?: string[] }) {
const texts = documents.map(({ pageContent }) => pageContent);
return this.addVectors(
await this.embeddings.embedDocuments(texts),
documents,
options
);
}
今回の事例だけ見ると無駄に見えますが（元々テキストを渡しているので）、Documentオブジェクトからテキスト情報だけを取り出して、OpenAI の APIを呼び出しembeddingをし、それを"addVectors()"に渡しています。おそらく処理の共通化の都合でこうなっているのだと思います。
this.embeddings.embedDocuments()はOpenAIのAPIを呼び出しembeddingを行なっています。
embeddingの処理はコード量が多いので引用を省きますので、上記URLで確認してください。
embedding実行時にパラメタを調整したい場合は fromTexts()を呼び出す際に渡す以下の引数を調整します。
new OpenAIEmbeddings()
LangChain側で初期値が設定されているので上記のようにそのまま渡すことが多いです。
LangChainはドキュメントも豊富で大いに助けられますが、ここら辺のパラメーター調整など、細かい部分はソースコードを直接読む方が速いです。公式ドキュメントはパラメタの設定方法についての言及が少なく、ドキュメントだけでは理解が難しいと感じることが多々あります。
そして、パラメタを調整しなくてもある程度動くので、通り過ぎてしまいがちですが、内部で設定されているパラメタに重要なものが結構あります。（思ったような動作をしない場合、ここに原因があることも多い）ソースコードは読めるようになった方が良いと考えています。
なお、embeddingのモデルの初期値は"text-embedding-ada-002"となっています。他に渡した複数のテキストをさらに分割して送信するなどのオプションがあります。
現状モデルが実質一択なので、技術ブログでここら辺の設定を変更している記事はあまり見かけません。
参考まで設定できるパラメタは以下のコードを参考にしてください。
export interface OpenAIEmbeddingsParams extends EmbeddingsParams {
/** Model name to use */
modelName: string;
/**
* Timeout to use when making requests to OpenAI.
*/
timeout?: number;
/**
* The maximum number of documents to embed in a single request. This is
* limited by the OpenAI API to a maximum of 2048.
*/
batchSize?: number;
/**
* Whether to strip new lines from the input text. This is recommended by
* OpenAI, but may not be suitable for all use cases.
*/
stripNewLines?: boolean;
}
話を戻して"addVectors"は受け付けた引数をSupabaseに送信する形に置き換えています。
async addVectors(
vectors: number[][],
documents: Document[],
options?: { ids?: string[] }
) {
const rows = vectors.map((embedding, idx) => ({
content: documents[idx].pageContent,
embedding,
metadata: documents[idx].metadata,
}));
// upsert returns 500/502/504 (yes really any of them) if given too many rows/characters
// ~2000 trips it, but my data is probably smaller than average pageContent and metadata
const chunkSize = 500;
let returnedIds: string[] = [];
for (let i = 0; i < rows.length; i += chunkSize) {
const chunk = rows.slice(i, i + chunkSize).map((row) => {
if (options?.ids) {
return { id: options.ids[i], ...row };
}
return row;
});
const res = await this.client.from(this.tableName).upsert(chunk).select();
if (res.error) {
throw new Error(
`Error inserting: ${res.error.message} ${res.status} ${res.statusText}`
);
}
if (res.data) {
returnedIds = returnedIds.concat(res.data.map((row) => row.id));
}
}
return returnedIds;
}
上記ではSupabaseのライブラリを使って、複数のレコードををupsertしています。
upsertといっても主キーを指定していないので、実質バルクインサート処理になっています。
また、ソースコードのコメントにあるように、一度にupsertする件数が多すぎるとエラーになるようで、500件を上限に分割して処理されているようです。コードの通り、この５００件の指定はパラメタで変更できません。
Supabaseのupsertについて詳しくは以下を確認してください。
Upsert data Javascript Client Library
パッケージのinstall、APIキーの設定
一旦、テキストをembeddingして、ベクトルデータベースに投入する部分を自前で実装をしていきます。
今回はMacのコマンドラインから実行するスクリプトを作ります。
以下のパッケージをinstallします
npm install @supabase/supabase-js openai dotenv
Supabaseのライブラリについては以下を
Javascript Client Library
OpenAIのライブラリは以下を参考にしてください。
openai-node
dotenvはAPIキーを .envファイルで管理するために使っています。
パッケージはKeyは以下のように読み込んでおきます。
ライブラリの初期化もしておきます。
import { createClient } from "@supabase/supabase-js";
import { Configuration, OpenAIApi } from "openai";
import dotenv from "dotenv";
dotenv.config();
if (!process.env.OPENAI_API_KEY)
throw new Error(`Expected env var OPENAI_API_KEY`);
if (!process.env.SUPABASE_ANON_KEY)
throw new Error(`Expected env var SUPABASE_ANON_KEY`);
if (!process.env.SUPABASE_URL) throw new Error(`Expected env var SUPABASE_URL`);
const client = createClient(
process.env.SUPABASE_URL,
process.env.SUPABASE_ANON_KEY
);
const configuration = new Configuration({
apiKey: process.env.OPENAI_API_KEY,
});
const openai = new OpenAIApi(configuration);
embeddings部分の実装
ベクトルデータベースに格納したい文章をembeddingします。
LangChainのサンプルコードだと味気ないので、ChatGPTに５つほど短い説明文を作成してもらいました。
文面の掲載は省きますが、以下のようなメタデータとともにベクトルデータベースに登録します。
{ name: "ベクトルDBについて", date: "20230721" },
{ name: "embeddingsについて", date: "20230721" },
{ name: "Node.jsについて", date: "20230721" },
{ name: "りんごについて", date: "20230721" },
{ name: "みかんについて", date: "20230721" },
OpenAIのライブラリ（Node.js)を使った、該当部分の実装は以下になります。
const response = await openai.createEmbedding({
model: "text-embedding-ada-002",
// 文章を配列で渡します。以下はダミーです。
input: ['ベクトルDBとはの本文(省略)', 'embeddingsとはの本文(省略)', 'Node.jsとはの本文(省略)' 'りんごについての本文(省略)', 'みかんについての本文(省略)'],
});
LangChainでは大量にembeddingしたときのタイムアウトやリトライ処理など実装してくれています。
今回は実装は省いてしまいましたが、こういった部分でも参考になる実装が多いです。
データの登録部分の実装
特筆するような実装はないのですが、500件の分割は省いてしまいました。
これは今後開発するサービス次第でケースバイケースで実装しましょう。。。
余談ですが、OpenAIのAPIの戻り値(今回はresponse)の処理は、構造が複雑でいつも間違えます。
let params = [];
for (let i = 0; i < response.data.data.length; i++) {
const data = {
content: texts[i],
embedding: response.data.data[i].embedding,
metadata: metadatas[i],
};
params.push(data);
}
const res = await client.from("documents").upsert(params).select();
if (res.error) {
throw new Error(
`Error inserting: ${res.error.message} ${res.status} ${res.statusText}`
);
}
console.log(res.data);
類似度検索のコードを読む
今度は投入したデータを検索する箇所のLangChain実装を読んでみます。
const resultOne = await vectorStore.similaritySearch("Hello world", 1);
similaritySearch()は文字列で検索クエリを受付け、それをembeddingしています。
引数に取得件数があり、サンプルコードでは１件が指定されていますが、通常類似度検索では結果を複数件取得して扱うことが多いです。（３件程度が多いですかね）
ベクトルデータベースや利用する検索ロジックにもよりますが、スコアが高いものが必ずしも一番有用とは限らず、かといって全てLLMに渡すのも合理的ではないので、数件LLMに渡して利用するパターンが多いのではないでしょうか。
ここの数字は実装時にサービスの内容を踏まえてチューニングすることになります。
検索ロジック、格納するテキストの内容や、格納する際の長さ、検索に用いられる文言などと組み合わせて調整します。（プロンプトの長さの制限も影響します）
検索処理はsimilaritySearchVectorWithScore()で行われます。
async similaritySearch(
query: string,
k = 4,
filter: this["FilterType"] | undefined = undefined
): Promise<Document[]> {
const results = await this.similaritySearchVectorWithScore(
await this.embeddings.embedQuery(query),
k,
filter
);
return results.map((result) => result[0]);
}
similaritySearchVectorWithScore()
この処理は受け付けた引数を使ってSupabaseのベクトルデータベースで作った関数を呼び出しています。
取得結果はLangChainでよく使われるDocumentオブジェクトに変換され、戻り値には類似度が付与されています。ただし、呼び出し元のsimilaritySearch()では類似度は使われません。
async similaritySearchVectorWithScore(
query: number[],
k: number,
filter?: this["FilterType"]
): Promise<[Document, number][]> {
if (filter && this.filter) {
throw new Error("cannot provide both `filter` and `this.filter`");
}
const _filter = filter ?? this.filter ?? {};
const matchDocumentsParams: Partial<SearchEmbeddingsParams> = {
query_embedding: query,
};
let filterFunction: SupabaseFilterRPCCall;
if (typeof _filter === "function") {
filterFunction = (rpcCall) => _filter(rpcCall).limit(k);
} else if (typeof _filter === "object") {
matchDocumentsParams.filter = _filter;
matchDocumentsParams.match_count = k;
filterFunction = (rpcCall) => rpcCall;
} else {
throw new Error("invalid filter type");
}
const rpcCall = this.client.rpc(this.queryName, matchDocumentsParams);
const { data: searches, error } = await filterFunction(rpcCall);
if (error) {
throw new Error(
`Error searching for documents: ${error.code} ${error.message} ${error.details}`
);
}
const result: [Document, number][] = (
searches as SearchEmbeddingsResponse[]
).map((resp) => [
new Document({
metadata: resp.metadata,
pageContent: resp.content,
}),
resp.similarity,
]);
return result;
}
Supabaseの関数の呼び出しについて、実装の説明時に扱います。
検索部分の実装
検索部分の実装です。
既に実装したembedding部分は省き、テーブル作成時、同時にSupabase上に作成した関数を呼び出します。
作成した関数の引数を用意し、関数名を指定して呼び出します。
LangChainではLangChain独自のデータ構造やさまざまな呼び出しパターンに対応するためコードが複雑になっていますが、メインの処理自体は拍子抜けするほどシンプルです。
詳細は以下を参考にしてください。
Call a Postgres function
// embedding
const text = "OpenAIの埋め込みについて教えてください。"
const response = await openai.createEmbedding({
model: "text-embedding-ada-002",
input: text,
});
// 関数に渡すパラメーター
const matchDocumentsParams = {
query_embedding: response.data.data[0].embedding,
match_count: 1
}
const { data, error } = await client.rpc(
"match_documents",
matchDocumentsParams
);
if (error) {
throw new Error(
`Error searching for documents: ${error.code} ${error.message} ${error.details}`
);
}
console.log(data)
検索を実行してみると、文言の一致がなくても、意味や文脈的に近いものが取れることがわかります。
例えば「果物」でりんごやみかんに関する記述、「システム開発」でNode.jsに関する記述は、RDSの通常の検索では取得できません。
今回のサンプルテキストは短く、件数も少ないので、是非様々な長文の検索を試してみてください。
メタデータでの絞り込み
せっかくなので、filterも使ってみましょう。
jsonb形式（jsonをバイナリ形式で保存）で格納されたデータで絞り込みを行えます。
上記コードのmatchDocumentsParamsを買い替えます。
const matchDocumentsParams = {
query_embedding: response.data.data[0].embedding,
filter: {name: "りんごについて"},
match_count: 1
}
上記の場合、nameに りんごについてが指定されているもののみマッチします。
今回登録したメタデータでマッチするのは
{ name: "りんごについて", date: "20230721" },
ですが、name以外の項目、「date」が指定されていなくてもマッチしました。
つまり保存されたjson全体がマッチする必要はありません。
指定した項目の中身が一致してればマッチします。
しかし、下記のように
filter: {name: "りんご"},
とすると、 nameが「りんご」で一致するレコードはないのでマッチしません。
今回テストデータで登録したレコード「りんごについて」マッチしません。
まとめ
今回はSupabaseのVectordDBにembeddingしたレコードを登録し、そのデータを類似度検索する処理を実装してみました。その際、LangChainのソースコードを参照にしつつも、LangChainは使わず実装しています。
OpenAIのAPIを使い始めて以来、LangChainを使って実装することがほとんどでしたが、最近はFunction callingをきっかけに、OpenAI APIのNode.jsライブラリを使って実装する機会も増えました。
LangChainで開発すると、サンプルが用意されていることもあり、定番の機能は素早く開発することができます。
何より、開発が活発で事例も多いためLangChainを通じて多くの知見を得ることができます。
一方で、こういう作業をしてみると、LangChainまかせになっているため、OpenAIのAPIや周辺の仕組みの理解がふわっとしたままになっていたことに気がつきます。LangChainはパッケージが大きいため、調査に時間がかりますし、更新頻度が多いため、アップデートなど保守コストなどデメリットも多少あるように思います。
月並みな結論になってしまいますが、適材適所で扱っていけるように、今後もLangChainとLangChainを使わない実装を併用していこうと思っています。
Discussion