LangChainのベクターストアとしてSupabaseを試してみた

Supabaseは、pgvectorを使用しているようで、他のPostgreSQL版（Amazon RDSなど）と組み合わせて使用することも可能です。Supabaseのプロジェクトを作成し、データベースを生成するSQLを実行します。また、Supabaseを使用してデータを挿入し、類似検索を行うAPIを作成することもできます。データの確認はTable EditorやSupabase SDKから行えます。ただし、RAGを使用する場合は負荷や影響を考慮する必要があります。

ベクターストアとしてsupabaseを試してました。内部的にはpgvectorを使っている様子なので、Amazon RDSなどのPostgreSQL版をOSS版Supabaseと組み合わせて使うとかもできるかもしれません。

SupabaseのDBを用意する

まずはdatabase.newからSupabaseのプロジェクトを新規作成します。

データベースを生成するSQLを実行します。DocsのSQLをSupabaseのダッシュボードで実行しましょう。

-- Enable the pgvector extension to work with embedding vectors
create extension vector;

-- 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
);

-- 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,
  embedding jsonb,
  similarity float
)
language plpgsql
as $
#variable_conflict use_column
begin
  return query
  select
    id,
    content,
    metadata,
    (embedding::text)::jsonb as embedding,
    1 - (documents.embedding <=> query_embedding) as similarity
  from documents
  where metadata @> filter
  order by documents.embedding <=> query_embedding
  limit match_count;
end;
$;

SQL Editorがあるので、GUIで実行できるのが便利ですね。

[Results]に成功メッセージが出ていれば、作成成功です。

Table Editorなどでリソースの作成が成功していることが確認できます。

HonoでAPIを作成する

ここからの実装は、Hono上で行います。ライブラリを追加でインストールしておきましょう。

npm install -S @supabase/supabase-js langchain

環境変数を3つ.dev.varsに設定します。

SUPABASE_PRIVATE_KEY=xxxx
SUPABASE_URL=https://xxxxx.supabase.co
OPENAI_API_KEY=sk-xxx

TypeScriptでのBindingsはこんな感じです。

import { Hono } from "hono";

export const supabaseApp = new Hono<{
    Bindings: {
        SUPABASE_PRIVATE_KEY: string;
        SUPABASE_URL: string;
        OPENAI_API_KEY: string;
    }
}>();

あとはインデックスへの投入を検索の処理を備えたAPIを用意するだけです。SupabaseVectorStore.fromTextsにデータを渡すと、EmbeddingされたデータがSupabaseにInsertされます。データに対する検索は、similaritySearchで行います。

import { SupabaseVectorStore } from "langchain/vectorstores/supabase";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { createClient } from "@supabase/supabase-js";

supabaseApp.get('ask', async c => {
    const privateKey = c.env.SUPABASE_PRIVATE_KEY;
    const url = c.env.SUPABASE_URL;
    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({
        openAIApiKey: c.env.OPENAI_API_KEY
      }),
      {
        client,
        tableName: "documents",
        queryName: "match_documents",
      }
    );
  
    const resultOne = await vectorStore.similaritySearch("Hello world", 1);
  
    console.log(resultOne);
    return c.json(resultOne)
})

実行すると、一致または類似するデータが返ってきました。similaritySearchの第二引数で件数を指定できます。

% curl http://127.0.0.1:8787/supabase/ask
[{"pageContent":"Hello world","metadata":{"id":2}}]

テーブルデータをSupabaseのTable Editorでみてみる

投入されたデータはSupabaseのGUIからも確認できます。

Supabase SDKで直接取得してみる

また、SupabaseのDBに投入されたデータなので、Supabase SDKからも取得できます。

supabaseApp.get('ask', async c => {
    const privateKey = c.env.SUPABASE_PRIVATE_KEY;
    const url = c.env.SUPABASE_URL;
    const client = createClient(url, privateKey);
    const data = await client.from('documents').select("id, content, metadata")
    return c.json(data)
});

通常のSQLを実行した形ですが、LangChainから投入したデータが見ることがわかります。

{
  "error": null,
  "data": [
    {
      "id": 1,
      "content": "Hello world",
      "metadata": {
        "id": 2
      }
    },
    {
      "id": 2,
      "content": "Bye bye",
      "metadata": {
        "id": 1
      }
    },
    {
      "id": 3,
      "content": "What's this?",
      "metadata": {
        "id": 3
      }
    }
  ],
  "count": null,
  "status": 200,
  "statusText": "OK"
}

もしEmbeddingしたベクトルデータが欲しい場合は、selectにembeddingを渡しましょう。

await client.from('documents').select("embedding")

ただしOpenAIやClaudeを使う場合、Embeddingした配列の長さが1,000を超えるため、必要になるまでは取らない方がデバッグがしやすいかなと思います。

[   {
      "embedding": "[0.0055157384,-0.009889425,0.016879791,-0.008615596,-0.030847976,0.00419171,-0.0043925107,-0.0003137508,-0.012154705,-0.02002985,0.017494744,-0.00610559,-0.0058389017,-0.011771929,0.014558036,0.006908792,0.03295638,0.0026417815,-0.0022637118,-0.021309953,-0.008866597,0.004263873,-0.003231633,-0.011577403,0.0031469204,0.024736112,0.017457092,-0.012913981,0.0033163456,-0.0123555055,0.03631979,-0.011458178,-2.2170905e-05,-0.012964182,-0.012757107,-0.009412523,0.0013938379,-0.04362391,0.0060240147,-0.0054341634,0.00015569883,0.008176345,-0.0013389314,-0.0026715877,-0.014357235,-0.0011083246,0.008948172,-0.032052778,0.0054843635,0.012556306,0.009431348,-0.00011883311,-0.026279764,-0.008609321,-0.022539856,0.004414473,-0.017996745,0.010924802,-0.015536939,0.008050845,-0.004389373,0.005569076,-0.023531308,0.013114782,0.016503291,0.001931136,-0.0056820265,-0.010027475,-0.027308868,-0.009832949,0.012719456,0.037800692,0.007680619,-0.009224272,0.014394886,-0.016992742,-0.045682114,-0.02068245,0.01580049,-0.0047721495,0.02813717,-0.017682994,-0.011301303,0.008
...

感想

実態がpgvectorらしいので、当然ではありますが、RDBMSでベクトルデータを扱えるのはいろいろと取り回しがききそうな予感がします。ただしRAGはAPI呼び出しが頻繁になる可能性がありますので、データベースへの負荷や通常のSQL呼び出し部分への影響などは考えた方がいい・・・のかもしれません。