Cloudflare VectorizeとLangChain＆OpenAIを組み合わせてRAGを作ってみた

Cloudflareのベクターストアを利用するためには、Workersを有料プランに変更する必要があります。月5ドルの基本料金で、リクエストが発生すればすぐにマネタイズできるため、気軽に利用できます。Vectorizeを使用してインデックスを作成し、Wranglerで設定値を追加します。また、VectorizeとOpenAIのAPIキーをHonoのBindingsに設定し、LangChainを使用して処理を書きます。最後に、Vectorizeのインデックスを削除する際にはWranglerコマンドを使用します。

Cloudflareのベクターストアを使ってみたかったので、LangChainとOpenAIのお馴染みセットで試してみました。

事前にWorkersを有料プランに変えないといけない

別の方が記事にされていましたが、Vectorizeを利用するにはWorkersを有料プランに変える必要があります。とはいえ基本料金は月５ドルで、従量課金が発生するほどのリクエストがくるならすぐにマネタイズした方がいいレベルなので、そこまで気負う必要はなさそうです。また、キャンセルしても1週間程度は利用できる様子でしたので、試してすぐキャンセルでもいいかもしれません。

Vectorizeでインデックスを作成する

Vectorizeでベクターストアのインデックスを作っていきましょう。Wrangler経由でもVectorizeは作成できます。作成時、dimensionsの値を指定する必要があります。

dimensionsの数字は、Embeddingに利用するモデルによって変わる様子でした。そのため、この記事で紹介していないモデルを利用する場合は、ドキュメントを必ず確認しましょう。768はWorkers AIで利用する@cf/baai/bge-base-en-v1.5ですが、OpenAIのada-002では1536を設定します。

npx wrangler vectorize create openai-rag-demo --dimensions=1536 --metric=cosine

作成に成功すると、wrangler.tomlに設定する値が表示されます。

📋 To start querying from a Worker, add the following binding configuration into
'wrangler.toml':

[[vectorize]]
binding = "VECTORIZE_INDEX" # available within your Worker on env.VECTORIZE_INDEX
index_name = "rag-demo"
✨  Done in 3.42s.

wrangler.tomlに項目を追加しましょう。この際、Workers AIの設定も追加しますので、次のような内容になるはずです。

name = "hono-langchain-demo"
compatibility_date = "2023-01-01"

[[vectorize]]
binding = "VECTORIZE_INDEX"
index_name = "openai-rag-demo"

Vectoriseを利用するアプリをローカル開発する場合、`--remote`オプションを設定しよう

wranglerを使ってローカル開発する場合、--remoteオプションを設定する必要があります。

npm run dev -- --remote

もしくはこのコマンドを利用します。

npx wrandler dev src/index.ts --remote

HonoのbindingsにVectoriseとOpenAI APIキーを設定する

作成したVectosizeのインデックスをWorkersのアプリと連携させましょう。Honoを利用するので、コンストラクタのGenericsに指定します。VectorizeIndexがそれですが、LangChainを利用する場合はほぼ内部的に使われるだけなので、anyでもあまり困りはしなさそうです。

import { Hono } from 'hono'
import { CloudflareVectorizeStore } from "langchain/vectorstores/cloudflare_vectorize";


const app = new Hono<{
    Bindings: {
        VECTORIZE_INDEX: VectorizeIndex;
    }
}>()

LangChainでOpenAI APIを利用するには、APIキーを取得する必要があります。ダッシュボードからAPIキーを取得しまておきましょう。そして取得した情報は、.dev.varsファイルにそれぞれ環境変数として設定します。

OPENAI_API_KEY=sk-xxxx

HonoのBindingsにも追加しましょう。

const app = new Hono<{
    Bindings: {
        VECTORIZE_INDEX: VectorizeIndex;
        OPENAI_API_KEY: string;
    }
}>()

LangChain.jsでOpenAI APIとVectorizeを利用した処理を書く

ここまでで下準備ができましたので、LangChainをインストールします。

% npm i langchain

EmbeddingとVector Storeの2つをインポートしましょう。

import { CloudflareVectorizeStore } from "langchain/vectorstores/cloudflare_vectorize";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";

それぞれをc.envから渡されるBindingsを使ってセットアップします。

app.post('/gen-index', async c => {
    const embeddings = new OpenAIEmbeddings({
        openAIApiKey: c.env.OPENAI_API_KEY
    })
    const store = new CloudflareVectorizeStore(embeddings, {
        index: c.env.VECTORIZE_INDEX
    })
    return c.json([])
})

続いてDBに投入するダミーデータを用意しましょう。

    const demoContents = [{
        id: 1,
        content: "Hello world"
    }, {
        id: 2,
        content: "good bye"
    }, {
        id: 3,
        content: "こんにちは"
}]

用意したダミーデータを、VectorizeStoreに投入するために成形します。

    const documents: Array<{
        pageContent: string;
        metadata: {
            postId: number;
        }
    }> = [];
    const documentIds: Array<string> = [];
    demoContents.forEach(content => {
        documents.push({
            pageContent: content.content,
            metadata: {
                postId: content.id,
            }
        });
        documentIds.push(content.id.toString());
    });

成形したデータをVectorizeStore経由でCloudflare Vectoriseに投入します。

　　await store.addDocuments(documents, { ids: documentIds });

これでDBのインデックスを作成するAPIができました。

app.post('/gen-index', async c => {
    const embeddings = new OpenAIEmbeddings({
        openAIApiKey: c.env.OPENAI_API_KEY
    })
    const store = new CloudflareVectorizeStore(embeddings, {
        index: c.env.VECTORIZE_INDEX
    })
    const demoContents = [{
        id: 1,
        content: "Hello world"
    }, {
        id: 2,
        content: "good bye"
    }, {
        id: 3,
        content: "こんにちは"
    }]
    const documents: Array<{
        pageContent: string;
        metadata: {
            postId: number;
        }
    }> = [];
    const documentIds: Array<string> = [];
    demoContents.forEach(content => {
        documents.push({
            pageContent: content.content,
            metadata: {
                postId: content.id,
            }
        });
        documentIds.push(content.id.toString());
    });
    await store.addDocuments(documents, { ids: documentIds });
    
    return new Response("updated", {
        status: 201
    })
})

このAPIをcurlで実行することで、Vectoriseにインデックスできます。

% curl http://127.0.0.1:8787/gen-index -XPOST
update

LangChainで、RAG APIを作る

Vectoriseにインデックスが作成できたので、いよいよ検索・RAGのAPIを作りましょう。

必要なクラスを初期化します。

import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { ChatOpenAI } from "langchain/chat_models/openai";
import { CloudflareVectorizeStore } from "langchain/vectorstores/cloudflare_vectorize";
import { StringOutputParser } from "langchain/schema/output_parser";
import { RunnableSequence } from "langchain/schema/runnable";
import {
  ChatPromptTemplate,
  HumanMessagePromptTemplate,
  AIMessagePromptTemplate,
} from "langchain/prompts";
import { formatDocumentsAsString } from "langchain/util/document";


...

app.get('/search', async c => {
    const model = new ChatOpenAI({
        temperature: 0,
        openAIApiKey: c.env.OPENAI_API_KEY,
        streaming: true,
        cache: true,
      });
    const embeddings = new OpenAIEmbeddings({
        openAIApiKey: c.env.OPENAI_API_KEY
    })
    const store = new CloudflareVectorizeStore(embeddings, {
        index: c.env.VECTORIZE_INDEX
    });
});

LangChainのドキュメントに記載されている、RAGのサンプルコードを追加しましょう。

app.get('/search', async c => {

    const model = new ChatOpenAI({
        temperature: 0,
        openAIApiKey: c.env.OPENAI_API_KEY,
        streaming: true,
        cache: true,
      });
    const embeddings = new OpenAIEmbeddings({
        openAIApiKey: c.env.OPENAI_API_KEY
    })
    const store = new CloudflareVectorizeStore(embeddings, {
        index: c.env.VECTORIZE_INDEX
    });

    const vectorStoreRetriever = store.asRetriever();

    const combineDocumentsPrompt = ChatPromptTemplate.fromMessages([
        AIMessagePromptTemplate.fromTemplate(
          "Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.\n\n{context}\n\n"
        ),
        HumanMessagePromptTemplate.fromTemplate("Question: {question}"),
      ]);
    const questionGeneratorTemplate = ChatPromptTemplate.fromMessages([
        AIMessagePromptTemplate.fromTemplate(
          "Given the following conversation about a codebase and a follow up question, rephrase the follow up question to be a standalone question."
        ),
        AIMessagePromptTemplate.fromTemplate(`Follow Up Input: {question}
      Standalone question:`),
      ]);
    const combineDocumentsChain = RunnableSequence.from([
        {
          question: (output: string) => output,
          context: async (output: string) => {
            const relevantDocs = await vectorStoreRetriever.getRelevantDocuments(output);
            return formatDocumentsAsString(relevantDocs);
          },
        },
        combineDocumentsPrompt,
        model,
        new StringOutputParser(),
      ]);
      
    const conversationalQaChain = RunnableSequence.from([
        {
          question: (i: { question: string }) => i.question,
        },
        questionGeneratorTemplate,
        model,
        new StringOutputParser(),
        combineDocumentsChain,
    ]);
    const question = "Tell me about Hono framework";
    const result = await conversationalQaChain.invoke({
        question,
    });
    return c.json(result)
})

このAPIを実行すると、LLMが生成したレスポンスを受け取ることができます。

curl -sN http://127.0.0.1:8787/search

後片付け: Vectorizeのインデックスを削除する

インデックスを削除する場合も、Wranglerコマンドから実行できます。

% npx wrangler vectorize delete rag-demo

warning package.json: No license field
$ /Users/okamotohidetaka/development/examples/hono-ai/node_modules/.bin/wrangler vectorize delete rag-demo
Deleting Vectorize index rag-demo
✔ OK to delete the index 'rag-demo'? … yes
✅ Deleted index rag-demo

一通り試してから、新しく作り直す用途などでも重宝しそうです。

参考

https://developers.cloudflare.com/vectorize/learning/create-indexes/