Supabase×Pythonで爆速RAG構築する

はじめに

みなさんは、RAGを構築するする際はどんなサービスを利用していますか？
代表的なものだと以下のようなクラウドサービスが挙げられると思います。

• AWS: knowledgebase×OpenSearch
• Azure: Azure AI Search, CosmosDB
• Google Cloud: Vertex AI Agent Builder
• Weaviate

ですが、個人で気軽にRAGを構築してアプリケーションをクラウドプラットフォームにデプロイするにはどれもお金がかかってしまい、なかなか手が出せないと思います。
ここで今回は、Supabaseを用いて無料で爆速RAG実装したので、その手法を紹介したいと思います！

Supabaseでは、500MBまで無料でデータベースを利用することができます

この記事の対象読者

• 無料でRAGを構築したい
• Supabaseを使ってみたい
• 個人開発で気軽にVectorDBを用いてRAGを構築したい
• LangChainを使わずにRAG構築したい

1. プロジェクトを作成する

こちら (https://supabase.com/dashboard/projects) からプロジェクトを作成してください。FREEプランの場合、2つまで無料でプロジェクトを作成することができます。
名前は適当で、リージョンは東京を選んでおけば良いと思います。

2. テーブルを作成する

次に、テーブルのセットアップをGUI上で行っていきます。SupabaseではPostgresDBをプロジェクトを作成してすぐ使うことができます。ここで、vectorカラムを扱えるように、以下のSQLをSQLEditorから打ち込んで、拡張機能を有効化します。

-- Example: enable the "vector" extension.
create extension vector
with
  schema extensions;

次に、テーブルスキーマを作成していきます。今回は、以下のスキーマを利用して説明を行なっていきます。ご自分のRAGの要件に合わせてスキーマをSQLEditorから打ち込んでください。

今回は、idと埋め込み対象のcontentのみを持つシンプルなスキーマを用います。また、embeddingはOpenAIの提供するembedding APIを用いるため、埋め込みの次元数は1536としています。こちらも、用いるembedding modelによって変える必要があります。(halfvecは最大4000次元までしかサポートしていません。)

また、今回は、コサイン距離（Cosine Distance）を用いて検索を行います。そのため、halfvec_cosine_opsという種類のindexを貼っています。これは、検索方法によって異なるため、適切な種類のindexを貼るようにしてください。
(参考: https://supabase.com/docs/guides/ai/vector-indexes/hnsw-indexes)

create table documents (
  id integer primary key generated always as identity,
  content text not null,
  embedding halfvec(1536),
);

create index on document using hnsw (embedding halfvec_cosine_ops);

3. クライアントを作成する

ここからPythonでコードを書いていきます。
まずは、DBクライアントを作成します。必要な認証情報は、SUPABASE_URLとSUPABASE_KEYの2つになります。こちら (https://supabase.com/dashboard/project/_/settings/api) から取得してください。
SUPABASE_KEYは2種類ありますが、今回は、権限の弱いpublicの方で大丈夫です。

また、Supabaseは2つのクライアント生成用メソッドを用意しています。本記事で紹介する範囲では、どちらでも実行することはできるのですが、今回は非同期クライアントを用います。

• create_client: 通常のDBクライアントを生成する
• create_async_client: 非同期DBクライアントを生成する

supabase_rag/client.py

from supabase import create_async_client, AsyncClient

class SupabaseClient():
    """SupabaseのDBクライアント"""

    def __init__(self):
        self._client = None

    async def get_client(self) -> AsyncClient:
        if self._client is None:
            self._client = await create_async_client(
                supabase_url=env.SUPABASE_URL,
                supabase_key=env.SUPABASE_KEY,
            )
        return self._client

4. データを挿入する

まずは、Pythonアプリケーション内で用いるモデルをpydanticを用いて記述します。特に型を意識しなければこの工程は不要です。

supabase/model.py

from pydantic import BaseModel

class Document(BaseModel):
    """
    Attributes:
        content(str): ドキュメントの内容
        embedding(list[float]): ドキュメントの埋め込み
    """

    content: str
    embedding: list[float]

class DocumentModel(Document):
    """
    supabaseのvector dbに格納されるデータ型

    Attributes:
        id: int
    """

    id: int

次に、埋め込み用のクライアントを作成します。

supabase_rag/embedding.py

from openai import OpenAI

class OpenAIEmbedding():
    """
    OpenAIの埋め込みクライアント
    """

    def __init__(self):
        self.client = OpenAI(api_key=env.OPENAI_API_KEY)

    def exec(self, text: str) -> list[float]:
        res = self.client.embeddings.create(input=text, model="text-embedding-3-small")
        return res.data[0].embedding

次に、Supabaseの挿入用クライアントを作成します。
記述にある通り、tableでテーブルを指定して、insertで挿入するデータを渡します。ここで、insertはdictを期待しているので、model_dumpを呼び出して、pydanticモデルからdictへの変換を行っています。
戻り値は、dataとcountが返ってきます。dataには挿入されたデータの配列が格納されているため、今回は、挿入したデータのidのみを返しています。

supabase_rag/insert.py

from supabase import AsyncClient

from supabase_rag.client import SupabaseClient
from supabase_rag.model import Document

class InsertSupabase():
    """
    supabaseのinsert用クライアント
    """

    def __init__(self, client: AsyncClient):
        self.client = client

    @classmethod
    async def new(cls, client: SupabaseClient):
        """ここからインスタンス作る"""
        conn = await client.get_client()
        return cls(conn)

    async def insert_document(
        self, data: Document, table_name: str = "documents"
    ) -> int:
        """
        supabaseにdocumentを挿入する

        Args:
            data(Document): 挿入するドキュメント
        Returns:
            int: 挿入したドキュメントのID
        """
        result = await self.client.table(table_name).insert(data.model_dump()).execute()
        return result.data[0]["id"]

最後に、作成したクライアントを統合したRAGクライアントを作成し、挿入操作を行なっていきます。

supabase_rag/rag.py

from supabase_rag.insert import InsertSupabase
from supabase_rag.embedding import OpenAIEmbedding
from supabase_rag.model import Document

class Rag():
    """
    RAGクライアント
    """

    def __init__(
        self,
        insert_client: InsertSupabase,
        embedding_client: OpenAIEmbedding
    ):
        self.insert_client = insert_client
        self.embedding_client = embedding_client

    async def insert_document(
        self, text: str
    ) -> int:
        embedding = self.embedding_client.exec(text)
        return await self.insert_client.insert_document(
            Document(
                content=text,
                embedding=embedding,
            )
        )

以下のスクリプトを実行して、SupabaseのGUIを見に行くとレコードが追加されていることが確認できると思います。

import asyncio

from supabase_rag.client import SupabaseClient
from supabase_rag.embedding import OpenAIEmbedding
from supabase_rag.insert import InsertSupabase

async def main():
    """デバッグ用"""
    client = SupabaseClient()
    insert_client = await InsertSupabase.new(client)
    embedding_client = OpenAIEmbedding()

    rag = Rag(insert_client, embedding_client)

    documents = [
        "Example Document 1",
        "Example Document 2",
        "Example Document 3",
        "Example Document 4",
        "Example Document 5",
    ]
    for doc in documents:
        document_id = await rag.insert_document(text=doc)
        print(f"Inserted ID : {document_id}")

if __name__ == "__main__":
    asyncio.run(main())

5. 検索を行う

では、待ちに待った検索を行なっていきましょう。
前述の通り、本記事のベクトル検索ではコサイン類似度を用いていきます。ですが、SupabaseのORMにはベクトル検索を行うためのメソッドは用意されていません。ここで、外部ライブラリに依存することなく実行する方法は2つあります。

• rpcメソッドを用いる
• execute_sqlメソッドを用いる

rpcメソッドは予め、SQLをSupabaseのSQLEditorから関数として定義しておいて、その関数を呼び出すという方法です。execute_sqlメソッドはその名の通り、SQLを直接実行できるメソッドになります。使い勝手は後者の方が良いのですが、クライアントのSUPABASE_KEYを権限が最も強いもの(Service Role Key)にする必要があります。
これは少し怖いので、今回はrpcメソッドを用いる方針でいきます。

まず、以下をSupabaseのSQLEditorから実行します。この関数はベクトルを引数に取り、コサイン距離で並び替えを行い、最も近い上位5つを返す関数になります。

/* semantic searchをするための関数 */
CREATE FUNCTION search_document(
    query VECTOR
) 
RETURNS setof documents
LANGUAGE SQL STABLE
AS $$
    SELECT *
    FROM documents
    ORDER BY documents.embedding <=> query ASC
    LIMIT 5;
$$;

では、登録した関数を用いて、検索を実装していきましょう。
呼び出し方は、第1引数に関数名、第2引数に辞書の形で引数を渡します。ここでも戻り値はdataに配列の形で格納されているので、それをDocumentModelに格納して返してあげましょう。

supabase_rag/search.py

from supabase import AsyncClient

from supabase_rag.client import SupabaseClient
from supabase_rag.model import DocumentModel

class SearchSupabase():
    """
    supabaseのsearch用クライアント
    """

    def __init__(self, client: AsyncClient):
        self.client = client

    @classmethod
    async def new(cls, client: SupabaseClient):
        """ここからインスタンス作る"""
        conn = await client.get_client()
        return cls(conn)

    async def search(self, query: list[float]) -> list[DocumentModel]:
        """
        DBからsemantic searchするメソッド
        """
        res = await self.client.rpc(
            "search_documents",
            {
                "query": query,
            },
        ).execute()
        return [DocumentModel(**item) for item in res.data]

最後に、検索用クライアントをRAGクライアントに組み込んでいきましょう。実装は簡単で、与えられた文字列を同じ埋め込み用クライアントで埋め込んであげて、先ほど作ったsearchメソッドを呼び出すだけです！

supabase_rag/rag.py

...
+from supabase_rag.search import SearchSupabase
-from supabase_rag.model import Document
+from supabase_rag.model import Document, DocumentModel

class Rag():
    """
    RAGクライアント
    """

    def __init__(
        self,
        insert_client: InsertSupabase,
+       search_client: SearchSupabase
        embedding_client: OpenAIEmbedding
    ):
        self.insert_client = insert_client
+       self.search_client = search_client
        self.embedding_client = embedding_client

    ...

+   async def search(self, query: str, category_name: str) -> list[DocumentModel]:
+       embedding = self.embedding_client.exec(query)
+       return await self.search_client.search(embedding)

以下は、検証用のスクリプトです。適当にドキュメントを挿入したのちに、適当なクエリを与えてテストしてみてください。良い感じに似た意味のドキュメントが5つ返ってくるはずです。

async def main2():
    client = SupabaseClient()
    insert_client = await InsertSupabase.new(client)
    search_client = await SearchSupabase.new(client)
    embedding_client = OpenAIEmbedding()

    rag = Rag(insert_client, search_client, embedding_client)

    result = await rag.search(
        query="Example search query",
    )
    for index, item in enumerate(result):
        print(f"result{index} : {item.content}")


if __name__ == "__main__":
    asyncio.run(main2())

補足

お気づきの方はいらっしゃると思いますが、かなり冗長にコードを書いています。実際のコードはそれぞれのクラスに対して基底クラスを作成し、それを継承する形で実装を行なっています。また、他のクラスにクラスを与えるときは、その基底クラスを型ヒントに用いています。

これは、RAGのサービスの変更に強くするためです。例えば今回はSupabaseを用いましたが、Weaviateに乗り換える場合は、全体を直すのはなく、supabase_rag/insert.pyとsupabase_rag/search.pyのみを修正すれば良いです。

記事用にもっと簡略化すれば良かったのですが、面倒でそのまま載せています😞
もちろん、ここまで冗長に書く必要はないので、SupabaseのRAGクライアントのクラスを1つにまとめて書いてあげても良いと思います。

supabase_rag/insert.py

from abc import ABCMeta, abstractmethod

# RAGクライアントのコンストラクタにはこちらのクラスを型ヒントとして与える
class Insert(metaclass=ABCMeta):
    @abstractmethod
    async def insert_document(
        self, data: Document, table_name: str = "documents"
    ) -> int:
        raise NotImplementedError()


class InsertSupabase(Insert):
    """
    supabaseのinsert用クライアント
    """

    def __init__(self, client: AsyncClient):
        ...

    async def insert_document(
        self, data: Document, table_name: str = "documents"
    ) -> int:
        ...

class InsertWeaviate(Insert):
    """
    weaviateのinsert用クライアント
    """

    def __init__(self, client: AsyncClient):
        ...

    async def insert_document(
        self, data: Document, table_name: str = "documents"
    ) -> int:
        ...

おわりに

最後まで読んでくださり、ありがとうございました！
個人で手を出すには敷居が高いと思われがちなRAGですが、Supabaseを用いることで簡単に、無料でRAG構築をすることができます。この機会にぜひ個人用のRAGを作ってみたり、簡単なアプリケーションに利用してみてください。

ご指摘やアドバイスがあればコメントで教えていただけると幸いです。