Cloudflare VectoriseとWorkersで関連記事検索

この記事は、「Cloudflare Advent Calendar 2023」3日目の記事です。

ベクトルDBとLLMを利用したRAGへの注目が一気に高まった2023年でした。自然言語ベースの検索や、回答文章の生成など、さまざまなユースケースでRAGの提案が進んでいますが、ネックになるのはLLMを利用する部分です。API呼び出し数やトークンによる課金だけでなく、DBと外部API呼び出しを行う関係から、レイテンシなども気になるケースもあります。

対策の1つとして、「ベクトルDBだけでできる機能はベクトルDBだけで実装すること」が考えられます。例えば関連記事のレコメンドであれば、「今いる記事のベクトルデータを利用して、類似性の高い記事を検索する」実装方法も可能です。

今回の記事では、Cloudflare Vectoriseを利用して関連記事検索を行う方法を紹介します。

Cloudflare Vectoriseは、Workersのenvからよびだす

Vectoriseの設定方法などについては割愛しますが、DBへのアクセスについては、R2やKVなどと同じくenvから行います。例えば今表示している記事のデータを取得する場合、await c.env.VECTORIZE_INDEX.getByIds([postId])で行います。

app.get('/:post_id/similar', async c => {
  const postId = c.req.param('post_id')
  const [postVector] = await c.env.VECTORIZE_INDEX.getByIds([postId])
  return c.json(postVector)
})

ポイントとしては、記事データをベクトルDBに投入する際、キーを記事IDに指定する必要があります。

記事のベクトルデータから、検索クエリを投げる

ベクトルDBでは、「ベクトルの距離が近いか遠いか」で検索を行うことができます。そのため、DBに保存されたベクトルデータと比較するためのベクトルデータがあれば、関連性の高いデータを検索できます。

今回のケースであれば、一度記事IDから該当記事のベクトルデータを取得後、そのデータを使って検索を行いましょう。

app.get('/:post_id/similar', async c => {
  const postId = c.req.param('post_id')
  const [postVector] = await c.env.VECTORIZE_INDEX.getByIds([postId])
  const similarPostsData = await c.env.VECTORIZE_INDEX.query(postVector.values, {})
  return c.json(similarPostsData)
})

レスポンスはこのようにインデックスのID( vectorId )と類似性( score )の2つが配列で返ってきます。

{
  "count": 5,
  "matches": [
    {
      "vectorId": "12992",
      "score": 1
    },
    {
      "vectorId": "12616",
      "score": 0.920572715
    },
    {
      "vectorId": "13156",
      "score": 0.907779246
    },
    {
      "vectorId": "13006",
      "score": 0.880790511
    },
    {
      "vectorId": "12589",
      "score": 0.859838598
    }
  ]
}

関連記事のタイトルや本文を取得する

正攻法であれば、ベクトルDBへの検索結果に含まれている記事IDを利用して、DBやREST APIへのデータ問い合わせを行います。

もう一つの方法としては、ベクトルDBのメタデータに記事タイトルやURLを保存し、それを利用することもできます。この方法の場合、IDの配列をクエリ結果から生成し、再びgetByIdsなどで取得を行います。

  const similarPostsData = await c.env.VECTORIZE_INDEX.query(postVector.values, {})
  const targetPostIds = similarPostsData.matches.map(match => {
    if (match.vectorId === postId) return null
    return match.vectorId
  }).filter((id): id is string => !!id)
  const postData = await c.env.VECTORIZE_INDEX.getByIds(targetPostIds)
  const similarPosts = postData.map(post => {
    return post.metadata
  })
  return c.json(similarPosts)

ベクトルDBのメタデータに本文を保存するのは、データの長さに制限があるため、あまり現実的ではありません。そのため、料金や速度と、どんなデータを取得したいかで実装方法を選ぶとよいでしょう。

まとめ

関連記事の取得を例に、Cloudflare VectoriseなどのベクトルDBを利用したAPIの実装方法を紹介しました。この API自体はLLMを利用しませんが、ベクトルDBに保存するデータの生成で、LLMのembedding APIなどを利用する必要があります。そのため、LangChainやCloudflare Workers AIなどをうまく活用しつつ、ベクトルDBの機能も把握していくことが重要となりそうです。

[Appendix]: 関連記事取得API全コード

app.get('/:post_id/similar', async c => {
  const postId = c.req.param('post_id')
  const [postVector] = await c.env.VECTORIZE_INDEX.getByIds([postId])
  const similarPostsData = await c.env.VECTORIZE_INDEX.query(postVector.values, {})
  const targetPostIds = similarPostsData.matches.map(match => {
    if (match.vectorId === postId) return null
    return match.vectorId
  }).filter((id): id is string => !!id)
  const postData = await c.env.VECTORIZE_INDEX.getByIds(targetPostIds)
  const similarPosts = postData.map(post => {
    return post.metadata
  })
  return c.json(similarPosts)
})