📖

RAGを作って学ぶCloudflareスタック

2024/08/08に公開

Cloudflare Workers

生成AIをよりよくする手法の一つにRAG = Retrieval-Augmented Generationがあります。これは単純な仕組みから作ることができて、効果的で面白いです。そして、Cloudflare Workersを中心としたCloudflareスタックで実現できます。やってみると、Cloudflareを使ったアプリケーション作成に必要なエッセンスをいくつも体験できることが分かりました。そこ今回は、シンプルなRAGアプリを作りつつ、Cloudflareスタックを学んでみましょう。

リポジトリ

今回扱うコードや関連する例は以下のリポジトリでみれます。

CloudflareスタックでRAGを作るとは？

RAGの実装方法を説明します。いくつかありますが、簡単な方法にします。

LLMと会話をするにはsystem、userというロールで以下のようなパラメータを渡します。

app.get('/', async (c) => {
  const result = await c.env.AI.run('@cf/meta/llama-3-8b-instruct', {
    messages: [
      { role: 'system', content: 'You are a good AI assistant' },
      { role: 'user', content: 'Hey AI!' }
    ]
  })
  return c.json(result)
})

RAGの場合はこのsystemのところに独自のコンテキストを渡します。例えば、Honoのことに詳しいAIにするには、質問に応じてHonoのWebサイトのコンテンツをとってきて、こんな感じにします。

const notes = [aboutBasicAuthForHono, aboutCORSForHono, aboutCSRCForHono]
const contextMessage = `Context: \n${notes.map((note) => `- ${note}`).join('\n')}`

app.get('/', async (c) => {
  const result = await c.env.AI.run('@cf/meta/llama-3-8b-instruct', {
    messages: [
      { role: 'system', context: contextMessage },
      { role: 'system', content: 'Answer the given question based on the context.' },
      { role: 'user', content: 'Hey AI!' }
    ]
  })
  return c.json(result)
})

コンテキストにテキストをそのまま渡しちゃうんですね。で、どんなテキストを渡すか？というと、質問の文に「近い」文章を渡します。上記の場合は「Honoについて」に近い文章をとってきてます。

ではこの「近さ」をどうやって計算するかというとCloudflareの場合はVectorizeというベクターデータベースを使います。また、コンテンツを保存しておくのにD1というSQLデータベースを使います。また、文章がどういうベクトルを持っているか？を計算するエンベッディングにはWorkers AIが使えます。

用意されたコンテンツを登録しておくための工程はこうなります。

コンテンツをいくつか用意する
それぞれのコンテンツのベクトルを出す（エンベッディング）
コンテンツをD1につくったテーブルのレコードに入れる
コンテンツのIDとベクトルをVectorizeに入れる

これで、各コンテンツのベクトルとIDがVectorizeに入ります。図にするとこうです。

次に、質問に返答するフェーズ。

質問をもらう
質問の文章をベクトルを出す
そのベクトルに近いベクトルをVectorizeから探す
取得したベクトルのコンテンツをD1からとってくる
コンテンツをコンテキスト用の文章にする
コンテキストと質問をLLMに渡して文章を生成する

この工程で我々は以下のCloudflareスタックを使うことになります。

Cloudflare Workers
D1
Vectorize
Workers AI - Embedding
Workers AI - Text Generation

また、今回はコンテンツを貯めておくストレージにD1を使っていますが、ここをKVやR2に変えることもできます。

では、いよいよRAGのアプリ作りたいところですが、上記した各要素を理解する必要があります。逆に言うとこれらを理解しておけば、組み合わせるだけでRAGができてしまいます。これは素晴らしい。

では、ひとつひとつ見ていきましょう。

1. 最初のWorker

まずは最初のCloudflare Workersのアプリを作ります。Create-Cloudflare CLI = C3というライブラリが提供されているのでをそれを使います。

npm create cloudflare@latest

テンプレートを聞かれるところがあるので「Hello World Worker」を選びましょう。あとは「TypeScript」で書くようにしましょう。

プロジェクトには以下のようなコマンドが用意されています。

開発サーバーを立ち上がる - npm run dev
Bindingsの型を生成する - npm run cf-typegen
デプロイする - npm run deploy

これらを使っていきます。

`fetch`ハンドラ

コードは以下のように書けば、JSONレスポンスを返せます。

src/index.ts

export default {
  async fetch() {
    return Response.json({
      message: 'Hello'
    })
  }
} satisfies ExportedHandler<Env>

fetchの中でResponseオブジェクトを返せばいいんですね。

`Request`

次に学んで欲しいのはリクエストのハンドリングです。例えば、リクエストURLを出力するにはこうすればいいです。

src/index.ts

export default {
  async fetch(req) {
    return Response.json({
      message: 'Hello',
      url: req.url
    })
  }
} satisfies ExportedHandler<Env>

fetchの第一引数がRequestオブジェクトなんですね。

環境変数

次に環境変数の扱いです。Clouflare Workersの場合ちょっと変わっています。wrangler.tomlにこう書きます。

wrangler.toml

[vars]
MY_VAR = "my-variable"

後述するBindingsの一つなんですが、とにかくwrangler.tomlに書きます。その後、TypeScriptの型を生成するために以下のコマンドを打ちます。

npm run cf-typegen

すると、env.MY_VARでその値にアクセスすることができます。

src/index.ts

export default {
  fetch: (req, env) => {
    return Response.json({
      message: 'Hello',
      url: req.url,
      myVar: env.MY_VAR
    })
  }
} satisfies ExportedHandler<{ MY_VAR: string }>

環境変数はwrangler.tomlに書かずともシークレットにすることもできます。

Secrets - Cloudflare docs

2. Honoを使う

素のfetchハンドラで開発してもいいのですが、ルーティングが大変だったりなので、Honoを入れましょう。

npm i hono

Honoは外部ライブラリに依存してなくhonoだけをインストールすればよいので、インストールは早いです。

コードを以下のように改造すれば、先程と同じことができます。

src/index.ts

import { Hono } from 'hono'

const app = new Hono<{ Bindings: Env }>()

app.get('/', (c) => {
  return c.json({
    message: 'Hello',
    path: c.req.path,
    myVar: c.env.MY_VAR
  })
})

export default app

ただし、このコードではルーティングがされています。つまり/にGETリクエストが来たときにハンドラが実行される。それ以外だと「404」を返すようになりました。

ミドルウェアを使う

Honoを使う強みはルーティング以外もあるのですが、豊富なミドルウェアが使えるのがその一つです。例えば、honoパッケージにビルトインのものは以下があります。

Basic Authentication
Bearer Authentication
Body Limit
Cache
Combine
Compress
CORS
CSRF Protection
ETag
IP Restriction
JSX Renderer
JWT
Logger
Method Override
Pretty JSON
Request ID
Secure Headers
Timeout
Timing
Trailing Slash

そのうちのPretty JSONを使ってみましょう。これは、URLに?prettyをつけるとJSONが整形されるというものです。使い方は簡単です。インポートして、適応したい場所にハンドラとして登録するだけです。

src/index.ts

import { Hono } from 'hono'
import { prettyJSON } from 'hono/pretty-json'

const app = new Hono<{ Bindings: Env }>()

app.use(prettyJSON())

app.get('/', (c) => {
  // ...
})

export default app

簡単ですね！

他の機能やより詳しいことについては公式サイトをご覧ください。

https://hono.dev/

3. D1を使う

次に、SQLデータベースのD1を使ってみます。

Bindings

D1やVectorize、それにKV、R2などのプロダクトをWorkersから扱うにはWeb APIを叩く方法以外に、Bindingsという方法があります。実は上記の環境変数の利用もBindingsと言えます。Bindingsはセキュアに快適な開発体験を維持したままリソースにアクセスできるよくできた考え方なので、この機会に触ってみましょう。

Bindingsを有効にするにはwrangler.tomlに設定を書き、プログラムないではenvもしくはHonoの場合はc.envからアクセスします。

D1の場合はまずデータベースをWranglerのCLIで作成します。

npm exec wrangler d1 create my-database

すると以下のような文字列がでます。

[[d1_databases]]
binding = "DB" # i.e. available in your Worker on env.DB
database_name = "my-database"
database_id = "your-db-id"

こいつをコピペして、wrangler.tomlに貼りましょう。これでDBの作成と設定ができました。次にテーブルを作りましょう。フィールドがidとtextだけのnotesを作ります。

npm exec wrangler d1 execute my-database -- --command "CREATE TABLE IF NOT EXISTS notes (id INTEGER PRIMARY KEY AUTOINCREMENT, text TEXT NOT NULL)"

では、コードからD1を操作してみましょう。最初に参照を書きます。GET /のリクエストを受けると、notesからコンテンツを取得して、ダンプするコードはこうなります。c.env.DBがD1のオブジェクトになっていて、そのままメソッドが実行できます。

src/index.ts

app.get('/', async (c) => {
  const result = await c.env.DB.prepare('SELECT * FROM notes').all()
  return c.json(result)
})

アクセスするとこんなJSONが返ります。resultsが空ですが、successがtrueなのでうまくいってますね！

{
  "success": true,
  "meta": {
    "served_by": "miniflare.db",
    "duration": 0,
    "changes": 0,
    "last_row_id": 0,
    "changed_db": false,
    "size_after": 16384,
    "rows_read": 1,
    "rows_written": 0
  },
  "results": []
}

更新を作りましょう。今回はPOST /notesにアクセスするとコンテンツがnotesに挿入されるようにします。コードは以下の通りです。いわゆるプリペアドステートメントを書いて、値をバインドしています。

src/index.ts

app.post('/notes', async (c) => {
  const defaultText = 'Today is a good day'
  const result = await c.env.DB.prepare('INSERT INTO notes (text) VALUES (?)')
    .bind(defaultText)
    .run()
  return c.json(result)
})

これだけだとdefaultTextの値しか挿入されないので、フォームから値を取れるようにしましょう。Honoの場合、c.req.parseBody()を使えます。

src/index.ts

app.post('/notes', async (c) => {
  const data = await c.req.parseBody<{ text: string }>()
  // ...
  const result = await c.env.DB.prepare('INSERT INTO notes (text) VALUES (?)')
    .bind(data.text ?? defaultText)
    .run()
  return c.json(result)
})

このコードを書いて、リクエストのフォームボディにtextというキーで値を入れるとそれが挿入されます。Postmanを使う場合はここに値を入れます。

これで簡単な更新と参照ができました。今回はSQLクエリを生で書いてましたが、クエリビルダやORMもあるので、試してみてください。

4. Workers AIを使う

Cloudflareのネットワーク上で、AIの推論が動くWorkers AIを使ってみましょう。

といっても、Workers AIの利用はとても簡単です。wrangler.tomlで以下の2行を追加するだけです。

wrangler.toml

[ai]
binding = "AI"

あとは、TypeScriptの型を追加するためにnpm run cf-typegenを実行しておきましょう。

ではコードを書いていきます。AIのオブジェクトにはc.env.AIでアクセスできます。次にrunと書くと、ずらっとAIモデルがIDEでサジェストされると思います。

Workers AIは利用できるモデルがたくさんあるのですね。どんなモデルがあるかは以下をみましょう。

Models - Cloudflare Docs

今回は言語生成にLlama 3を使います。propmt引数に言葉を渡すと返答が返ってきます。超簡単に書ける！

src/index.ts

app.get('/ai', async (c) => {
  const result = await c.env.AI.run('@cf/meta/llama-3-8b-instruct', {
    prompt: 'Hey AI',
  })
  return c.json(result)
})

これだけだと、同じ質問だけになってしまうので、クエリパラムを受け取るようにしましょう。Honoだとc.req.query()でアクセスできます。

src/index.ts

app.get('/ai', async (c) => {
  const prompt = c.req.query('prompt')
  const result = await c.env.AI.run('@cf/meta/llama-3-8b-instruct', {
    prompt: prompt ?? 'Hey AI',
  })
  return c.json(result)
})

5. Vectorizeを使う

VectorizeはCloudflareが持つベクターデータベースです。ベクターデータベースとはドクトル群を挿入することができ、あるベクトルに近いベクトルを出してくれます。今回はRAGのために使っていますが、他にも分類、レコメンドなどに使うことができます。また、Cloudflareスタックと組み合わせるのではなく、Vectorize単体で利用もできます。

興味深い例として、StripeのAPIと組み合わせてEコマースサイトのレコメンドを作るチュートリアルがCloudflareのドキュメントにあります。

Recommend products on e-commerce sites using Workers AI and Stripe - Cloudflare docs

wranglerコマンドでVectorizeのデータベースを作ってみましょう。以下はディメンションが3のデータベースを作っています。

npm exec wrangler vectorize create my-index-vectorize -- --dimensions=3 --metric=cosine

できたら、他のBindingsと同じようにwrangler.tomlに設定を書きます。

wrangler.toml

[[vectorize]]
binding = "VECTORIZE_INDEX"
index_name = "my-index-vectorize"

npm run cf-typegenも実行しておきましょう。

まずはVectorizeにベクトルを挿入するコードを書きます。サンプルとしてハードコードしたデータを使います。ディメンションが3なので、valuesの長さが3になります。また、Vectorizeには各値にメタデータを追加できるので、今回は商品URLのような値を入れておきます。

src/index.ts

// Based on https://developers.cloudflare.com/vectorize/get-started/intro/#4-insert-vectorsd
const sampleVectors: Array<VectorizeVector> = [
  {
    id: '1',
    values: [32.4, 74.1, 3.2],
    metadata: { url: '/products/sku/13913913' },
  },
  {
    id: '2',
    values: [15.1, 19.2, 15.8],
    metadata: { url: '/products/sku/10148191' },
  },
  {
    id: '3',
    values: [0.16, 1.2, 3.8],
    metadata: { url: '/products/sku/97913813' },
  },
  {
    id: '4',
    values: [75.1, 67.1, 29.9],
    metadata: { url: '/products/sku/418313' },
  },
  {
    id: '5',
    values: [58.8, 6.7, 3.4],
    metadata: { url: '/products/sku/55519183' },
  },
]

ハンドラは以下の通りです。これも簡単です。

src/index.ts

app.get('/vectorize/insert', async (c) => {
  const inserted = await c.env.VECTORIZE_INDEX.insert(sampleVectors)
  return c.json(inserted)
})

さて、GET /vectorize/insertにアクセスしてみましょう。するとInternal Server Errorが出ます。これは想定内です。Vectorizeはリモートで実行する必要があるので、これまでのローカル前提の開発サーバーではエラーがでるのです。そこで、package.jsonを編集して、wrangler devコマンドに--remote引数を渡すようにしましょう。

package.json

{
  "scripts": {
    "dev": "wrangler dev --remote"
   }
}

これで、もう一度開発サーバーを立ち上げなおすとVectorizeへの挿入がうまくいくはずです。

取得は対象のベクトルをqueryメソッドに渡します。topKに値を指定しています。これは、上位何件のベクトルを取得するか？を指定するもので、今回は上位2個をとってきています。

src/index.ts

app.get('/vectorize/query', async (c) => {
  const queryVector: Array<number> = [54.8, 5.5, 3.1]
  const matches = await c.env.VECTORIZE_INDEX.query(queryVector, {
    topK: 2,
    returnValues: true,
    returnMetadata: true,
  })
  return c.json(matches)
})

これでベクトルデータベースも使うことができました。

6. Embeddingを作る

RAGを実現するには、コンテンツの文章がどのようなベクトルを持っているかを解析して、それをVectorizeに入れたり、クエリにしなくてはいけません。このベクトルをテキストから出すのがこの工程です。といっても、Workers AIにそのためのモデルがあります。

このコードではテキストエンベッディングのための@cf/baai/bge-base-en-v1.5というモデルを叩いています。

src/index.ts

app.get('/embedding', async (c) => {
  const text = 'Iekei ramen is one of the Ramen category'
  const result = await c.env.AI.run('@cf/baai/bge-base-en-v1.5', {
    text,
  })
  return c.json(result)
})

実行すると768ディメンションのベクトルが取れるはずです。

{
    "shape": [
        1,
        768
    ],
    "data": [
        [
            0.021004894748330116,
            0.008485095575451851,
            -0.03035151958465576,
            0.013451575301587582,
            0.0438392348587513,

7. RAGアプリケーションの構築

以上で、はじめてのWorkerから、Honoの導入。そして以下のBindingsを試してきました。

D1
Vectorize
Workers AI - Embedding
Workers AI - Text Generation

これだけ揃えばRAGを作ることができるのです！

Bindingsはこれまで使ったものを利用できます。2点だけやることがあります。

Vectorizeのディメンションを変更する必要があるので、一旦削除して、ディメンション768で作り直しましょう。

npm exec wrangler vectorize delete my-index-vectorize
npm exec wrangler vectorize create my-index-vectorize -- --dimensions=768 --metric=cosin

D1のリモートにはnotesテーブルが存在しないので、--remoteオプションをつけてつくります。

npm exec wrangler d1 execute my-database -- --remote --command "CREATE TABLE IF NOT EXISTS notes (id INTEGER PRIMARY KEY AUTOINCREMENT, text TEXT NOT NULL)"

さて、あとはこれまでのコードを組み合わせるだけ！コンテンツを用意するエンドポイントPOST /notesと質問に答えるGET /を書いてみましょう。少々長いですが、コメント付きで掲載します。

コンテンツを用意する

src/index.ts

app.post('/notes', async (c) => {
  // フォームボディから`text`を取得
  const { text } = await c.req.parseBody<{ text: string }>()
  // textがなければ`/`へリダイレクト
  if (!text) {
    return c.redirect('/')
  }

  // コンテンツをD1に挿入する
  const d1Result = await c.env.DB.prepare('INSERT INTO notes (text) VALUES (?)')
    .bind(text)
    .run()
  if (!d1Result.success) {
    return c.text('Fail to create a note', 500)
  }

  // コンテンツのベクトルを取得
  const aiResult = await c.env.AI.run('@cf/baai/bge-base-en-v1.5', {
    text,
  })
  const values = aiResult.data[0]
  if (!values) {
    return c.text('Fail to generate vector embeddings')
  }

  // D1に入れた際のIDとベクトルをVectorizeへ入れる
  const vectorizeResult = await c.env.VECTORIZE_INDEX.insert([
    {
      id: d1Result.meta.last_row_id.toString(),
      values,
    },
  ])

  return c.json(vectorizeResult)
})

質問に答える

src/index.ts

app.get('/', async (c) => {
  // クエリパラメータから質問の文章を取得
  const text = c.req.query('text') || 'What is the square root of 9?'

  // 質問のベクトルを取得
  const embeddingResult = await c.env.AI.run('@cf/baai/bge-base-en-v1.5', {
    text,
  })

  // 質問のベクトルに近いベクトルを上位5件持ってくる
  const vectorizeResult = await c.env.VECTORIZE_INDEX.query(
    embeddingResult.data[0],
    {
      topK: 5,
    }
  )

  // スコアが0.7より上の値ではないと採用しない
  const CUTOFF = 0.7
  // ベクトルのID群をつくる
  const vectorIds = vectorizeResult.matches
    .filter((vector) => {
      return vector.score > CUTOFF
    })
    .map((vector) => vector.id)

  const notes: string[] = []

  if (vectorIds.length) {
    // ベクトルのID群のコンテンツを取得する
    const query = `SELECT * FROM notes WHERE id IN (${vectorIds.join(', ')})`
    const d1Result = await c.env.DB.prepare(query).bind().all()
    if (d1Result.success) {
      d1Result.results.forEach((entry) => {
        notes.push(entry.text as string)
      })
    }
  }

  // 取得したベクトルのコンテツを文字列連結してコンテツにする
  const contextMessage = notes.length
    ? `Context: \n${notes.map((note) => `- ${note}`).join('\n')}`
    : ''
  const systemPrompt = 'Answer the given question based on the context.'
  const messages: RoleScopedChatInput[] = [
    { role: 'system', content: systemPrompt },
    { role: 'user', content: text },
  ]

  // コンテキストをメッセージの冒頭に入れる
  if (contextMessage) {
    messages.unshift({ role: 'system', content: contextMessage })
  }

  // テキスト生成を実行する
  const aiResult = await c.env.AI.run('@cf/meta/llama-3-8b-instruct', {
    messages,
  })

  return c.text(aiResult.response)
})

試す

実際に今回つくったRAGアプリを試してみました。上記した通り、残念ながら日本語に弱いので、今回はHonoのWebサイトに掲載されているドキュメントを挿入してみました。あまりドキュメントが長いとテキスト生成のAIに食わせるときに文字数がオーバーするので、ある程度カットします。

この状態で"How to implement authentications on Hono?"と尋ねるとしっかりとHonoらしいコードでBasic認証、Bearer認証、JWT認証、カスタム認証があると返してくれました。

これはいい感じです！英語かつテクニカルドキュメントなら期待した挙動をしそうですね。

扱わなかったこと

今回扱いたかったけど扱わなかったことがあります。

Function Calling
マルチモーダル
KVを使う
R2を使う

これは今後紹介するかもしれませんが、Cloudflareのドキュメントが詳しいのでそれをご覧ください！

https://developers.cloudflare.com/

まとめ

以上、CloudflareスタックでRAGアプリケーションをつくってみました。その過程で以下を使う方法を学びました。

Cloudflare Workers
D1
Vectorize
Workers AI - Embedding
Workers AI - Text Generation

RAGは面白いし、これらの技術を個別に使う、違う組み合わせで使うこともできます。そしてあなたなりのアプリケーションを作ってください！

参考文献

以下を参考にしました。

リポジトリ

CloudflareスタックでRAGを作るとは？

1. 最初のWorker

fetchハンドラ

Request

環境変数

2. Honoを使う

ミドルウェアを使う

3. D1を使う

Bindings

4. Workers AIを使う

5. Vectorizeを使う

6. Embeddingを作る

7. RAGアプリケーションの構築

コンテンツを用意する

質問に答える

試す

扱わなかったこと

まとめ

参考文献

Discussion

`fetch`ハンドラ

`Request`