こんばんは。

はじめに

旅行行きたいなー、でも行く場所決めるの大変だなぁ、とか思っていまして、時代はAIですのでChatGPTをつかって旅行先をリコメンドしてもらっていました。

そこからしばらくして、自分でもおすすめの旅行先をリコメンドするエンジンを書けないかなと思案していて、キーワードからおすすめの旅行先を出す、というものを作っていました。

当初はリコメンドする部分については検索キーワード分割をしたり、付近のおすすめ施設をGoogleMapAPIをつかってDBに格納しつつ重み付けして優先度を…としていたのですが、なかなか大変でした。

色々調べていてPineconeというベクターDBを使って検索させる方法を使えばキーワードの関連度からリコメンドしてくれそうでした。

などと思っていたら、SupabaseからSupabase VecotorなるベクターDBが提供されていて、Pineconeからそれに乗り換えた、という記事を偶然Xで発見しました。
https://supabase.com/customers/mendableai

これは便利そう！ということで早速使ってみました。

方針

リコメンドの方針について下記のようにすることにしました。

OpenAIのChatGPTを使い、旅行にリコメンドする施設情報をベクター化する
ベクター化された施設をDBに保存する
検索キーワードから近似値検索を行い、施設を取り出す

この保存するDBにSupabase Vecotorを使いました。

環境整備

Supabaseはローカル環境が用意されているのでざっくり準備しました。
マイグレーションを作成してテーブルなどを準備します。

supabase migration new load_vector_extension

pgvectorの準備

PostgreSQLで提供されるベクターDBを有効化します。

create extension vector; -- ここでSupabase Vectorを使う宣言

次にリコメンドしたい施設を保存するテーブルを作成します。

create table documents (
  id bigserial primary key,
  content text,
  embedding vector(1536) -- ここで1536次元のベクター型にします
);

検索する関数をPostgreSQL側で定義します。

-- ドキュメントを検索する関数
-- query_embedding: 検索クエリの埋め込みベクトル (vector(1536))
-- match_threshold: 類似度の閾値 (float)
-- match_count: 返されるドキュメントの最大数 (int)

create or replace function match_documents (
  query_embedding vector(1536),
  match_threshold float,
  match_count int
)
-- 返り値:
-- id: ドキュメントのID (bigint)
-- content: ドキュメントの内容 (text)
-- similarity: ドキュメントとクエリの類似度 (float)
returns table (
  id bigint,
  content text,
  similarity float
)
language sql stable
as $$
  select
    documents.id,
    documents.content,
    1 - (documents.embedding <=> query_embedding) as similarity
  from documents
  where documents.embedding <=> query_embedding < 1 - match_threshold
  order by documents.embedding <=> query_embedding
  limit match_count;
$$;

最後にIVFFlatのインデックスを作成しておきます。

create index on documents using ivfflat (embedding vector_cosine_ops)
with
  (lists = 100);

コード

まずは施設を保存するコードを書きます。
OpenAIで施設名をEmbeddingし保存します。
OpenAIのAPIでは少し前にtext-embedding-3-smallという新しいモデルが利用可能になっていて、このモデルを使用します。次元は1536です。

TypeScriptで下記のようなコードでEmbeddingされます。便利ですね！

  const openai = new OpenAI({ apiKey: OPENAI_API_KEY })
  const embeddingResponse = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: document,
  })
  const embedding = embeddingResponse.data[0].embedding

これをAPIとして呼び出せるようにしたものです。

functions/search-document/index.ts


/* To invoke locally:  
  curl -i --location --request POST 'http://127.0.0.1:54321/functions/v1/search-document' \
    --header 'Authorization: Bearer '
    --header 'Content-Type: application/json' \
    --data '{"query":"Functions"}'
*/

import "jsr:@supabase/functions-js/edge-runtime.d.ts"
import OpenAI from 'https://deno.land/x/openai@v4.24.0/mod.ts'
import { createClient } from 'jsr:@supabase/supabase-js@2'

const SUPABASE_URL = Deno.env.get('SUPABASE_URL')
const SUPABASE_ANON_KEY = Deno.env.get('SUPABASE_ANON_KEY')
const OPENAI_API_KEY = Deno.env.get('OPENAI_API_KEY')

Deno.serve(async (req) => {
  
  if (!SUPABASE_URL || !SUPABASE_ANON_KEY || !OPENAI_API_KEY) {
    return new Response(
      JSON.stringify({ error: 'Environment variables are not set properly' }),
      { status: 500 }
    )
  }
  
  const supabaseClient = createClient(SUPABASE_URL, SUPABASE_ANON_KEY)

  const { query } = await req.json()
  const document = query.replace(/\n/g, ' ')

  const openai = new OpenAI({ apiKey: OPENAI_API_KEY })
  const embeddingResponse = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: document,
  })
  const embedding = embeddingResponse.data[0].embedding

  const { data, error } = await supabaseClient.rpc('match_documents', {
    query_embedding: embedding,
    match_threshold: 0.5,
    match_count: 5,
  })

  if (error) {
    return new Response(
      JSON.stringify({ error: error.message }),
      { status: 500 }
    )
  }

  return new Response(
    JSON.stringify({ data }),
    { status: 200, headers: { 'Content-Type': 'application/json' } }
  )
})

実際に呼び出してみます。

Supabase Studioで確認します。

保存されているようです！（Studio便利ですね！）

検索

検索も同様に、検索キーワードをEmbeddingして検索します。検索は先程の関数を利用します。

functions/search-document/index.ts

/* To invoke locally:  
  curl -i --location --request POST 'http://127.0.0.1:54321/functions/v1/search-document' \
    --header 'Authorization: Bearer ' \
    --header 'Content-Type: application/json' \
    --data '{"query":"Functions"}'
*/

import "jsr:@supabase/functions-js/edge-runtime.d.ts"
import OpenAI from 'https://deno.land/x/openai@v4.24.0/mod.ts'
import { createClient } from 'jsr:@supabase/supabase-js@2'

const SUPABASE_URL = Deno.env.get('SUPABASE_URL')
const SUPABASE_ANON_KEY = Deno.env.get('SUPABASE_ANON_KEY')
const OPENAI_API_KEY = Deno.env.get('OPENAI_API_KEY')

Deno.serve(async (req) => {
  
  if (!SUPABASE_URL || !SUPABASE_ANON_KEY || !OPENAI_API_KEY) {
    return new Response(
      JSON.stringify({ error: 'Environment variables are not set properly' }),
      { status: 500 }
    )
  }
  
  const supabaseClient = createClient(SUPABASE_URL, SUPABASE_ANON_KEY)

  const { query } = await req.json()
  const document = query.replace(/\n/g, ' ')

  const openai = new OpenAI({ apiKey: OPENAI_API_KEY })
  const embeddingResponse = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: document,
  })
  const embedding = embeddingResponse.data[0].embedding

  const { data, error } = await supabaseClient.rpc('match_documents', {
    query_embedding: embedding,
    match_threshold: 0.0,
    match_count: 5,
  })

  if (error) {
    return new Response(
      JSON.stringify({ error: error.message }),
      { status: 500 }
    )
  }

  return new Response(
    JSON.stringify({ data }),
    { status: 200, headers: { 'Content-Type': 'application/json' } }
  )
})

検索してみます。

なぜかうまくヒットしません…。

関数周りなどの調整も調整したところなんとか取れました。

使ってみて

思った以上に簡単にベクターDBを使うことができましたが、検索結果は全然うまく取れませんでした。

Supabaseで提供されていたので、アプリ全体のバックエンドをSupabaseにまかせて後はフロントを実装すれば良さそうでそこはすごく便利でした。