🐡

Cloudflare Vectorizeで日本語検索

2024/02/10に公開

Cloudflare

OpenAI

RAG

tech

Cloudflare VectorizeはCloudflareがホストするVector database

PineconeのようにHTTP経由で呼び出して使う

Workers AIと組合せてllama-2とかでRAGを作ってくれという想定らしいけどホストされているText Embeddingsのモデルが英語用しかない

埋め込み表現に変換してVector databaseのAPIに投げるだけなので保存するデータはどのモデルを使っても問題はないのだけど、検索をする時にCloudflare Workersから使いたかったのでHTTP呼び出し可能なものにする

今回はOpenAIのtext-embedding-3の新モデルを試すことにした

サンプルデータを登録してクエリで牽くという段階までは以下のドキュメントどうりに実行すると実現できるので省略する

今回は日本語検索をしたくてOpenAIのtext-embedding-3を使いたいので、それにあわせたインデックスを作る

npx wrangler vectorize create my-index --dimensions=1536 --metric=cosine

--dimensions=1536は By default, the length of the embedding vector will be 1536 for text-embedding-3-small or 3072 for text-embedding-3-large. を参考にした(あと3072に指定したら内部エラーが起きた)

料金は https://openai.com/pricing

トークン数の目安は https://platform.openai.com/tokenizer で調べられるがtext-embedding-3でも同じ結果になるのかは調べてない

For third-generation embedding models like text-embedding-3-small, use the cl100k_base encoding.という記述はある

--metric=cosineは以下の3種類から選べそうだけどドキュメントに沿ってcosine(コサイン類似度)にした

データ登録API

GET STARTEDではハードコードしたベクトルデータをINSERTしているだけなので、ここを外部から入力できるようにする

let path = new URL(request.url).pathname;
if (path.startsWith("/insert")) {
    // { id: number, values: number[], metadata: {} },
    const payload: Array<VectorizeVector> = await request.json();
    let inserted = await env.VECTORIZE_INDEX.insert(payload);
    return Response.json(inserted);
}

idがインデックスの一意なIDで、valuesが[32.4, 74.1, 3.2]などの埋め込み表現(embedding)でmetadataが自由に保存できるデータになる

これで wrangler dev --remote で起動して POST http://localhost:8080/insert にデータを送るとインデックスに登録できるようにはなった

ただこれをデプロイすると第三者が勝手に登録できてしまうので簡易認証を追加するなりローカルだけで動かすなりしないといけない

Workers上でINSERTする以外の選択肢としてはInsert VectorsのREST APIが使える

このREST APIはwranglerに組込まれているので

wrangler vectorize insert --file data.ndjson

から登録できる。.ndjsonはjsonl形式のような1行に有効なjsonが入っているテキストファイル

今回やったことは

検索対象のテキストを用意して分割してから
openaiのnpmモジュールでembeddingに変換
{ id: number, values: number[], metadata: {} } の形式にしてPOST http://localhost:8080/insert

1-2の単純なテキスト処理の部分は省略する

クエリ検索API

Workersで検索結果を返すAPIを作る

ユーザーの入力をtext-embedding-3-smallに変換
env.VECTORIZE_INDEX.query()に渡す
結果をレスポンスで返す

(htmxは使ってみたかっただけで、普通にJSONレスポンスでもいい)

if (path === "/query" && request.method === "POST") {
    const formDataText = await request.text();
    const params = new URLSearchParams(formDataText);
    const input = params.get("query") ?? '';
    if (input.length === 0) {
        return new Response("Bad Request", { status: 400 });
    }

    const openai = new OpenAI({ apiKey: env.OPENAI_API_KEY });
    const openaiResponse = await openai.embeddings.create({ model: "text-embedding-3-small", input });
    const queryVector = openaiResponse.data[0].embedding;
    console.log('queryVector: ', queryVector.length);

    const result = await env.VECTORIZE_INDEX.query(queryVector, { topK: 10, returnValues: true, returnMetadata: true });

    const htmlResponse = `
        ${result.matches.map(match => {
        const chunkString = typeof match.metadata?.chunk === 'string'
            ? match.metadata.chunk.slice(0, 128) + '...'
            : match.metadata?.chunk;
        return `
            <tr>
                <td>${match.id}</td>
                <td>
                    <a href="https://digital-gov.note.jp/n/${match.metadata?.key}">
                    ${match.metadata?.name}
                    </a>
                </td>
                <td>${chunkString}</td>
                <td>${match.score}</td>
            </tr>
            `
    }).join('')}
    `;

    return new Response(htmlResponse, { headers: { "Content-Type": "text/html" } });
}

env.OPENAI_API_KEYはwrangler secret put OPENAI_API_KEYで設定できる。以下のチュートリアルなどが参考になる

検索実行画面

検索画面は以下のように用意した(手抜きフルセットみたいなデッキ)

if (path === "/" && request.method === "GET") {
    const htmlContent = `
    <html>
        <head>
            <meta charset="UTF-8">
            <title>VectorizeIndex Demo</title>
            <meta name="viewport" content="width=device-width, initial-scale=1">
            <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/water.css@2/out/water.css">
        </head>
        <body>
            <h1>VectorizeIndex Demo</h1>
            <p><a href="https://digital-gov.note.jp/">デジタル庁公式note</a>の記事を<a href="https://developers.cloudflare.com/vectorize/">Cloudflare Vectorize</a>を使ってコサイン類似度に基づいて検索します。</p>
            <form hx-post="/query" hx-target="#results" method="POST">
                    <textarea type="text" name="query">ガバメントクラウドとは何ですか?</textarea>
                    <button type="submit">Ask</button>
            </form>
            <table>
            <thead>
                <tr>
                    <th>ID</th>
                    <th>Article</th>
                    <th>Chunk</th>
                    <th>Score</th>
                </tr>
            </thead>
            <tbody id="results"></tbody>
            </table>
            <script src="https://unpkg.com/htmx.org"></script>
        </body>
    </html>
`;
    const headers = { "Content-Type": "text/html" };
    return new Response(htmlContent, { headers });
}

実際に動かしてみた例が以下。ブログをソースにチャンクを拾ってきてRAGのバックエンドにする部分だけ作った

DEMO: https://vectorize.laiso.workers.dev/

Discussion

料金はクエリとインデックスで課金され、インデックス量と実行回数にdimensionsをかける方式と「Pricing · Vectorize」に書いてあった
OpenAIのAPI叩いてるからエッジlocation→USのレイテンシが発生しているはず。AI GatewayのキャッシュはあるがRAG用途には向かなさそう