PGliteとClaude APIをつかってクライアントだけでRAGする

中村修太

2024.09.02

Introduction

ここにあるように、Anthropic社がClaude APIに対してCORSサポートを追加しました。
これによってブラウザから直接Anthropicのモデルを呼び出すことができるようになってます。

ちょっと前にPGliteというWASM版PostgreSQLもリリースされました。
これはpgvector(Postgres用のvector類似性検索)などの拡張機能も使えます。

今回はこれらを使用してクライアントサイドのみでVector検索とClaude apiを使って
RAGをやってみます。

Claude API with CORS

Claude APIがCORSサポートされたことで、
クライアントサイドから直接Claude APIが呼び出せるようになります。
いままではサーバを立ててそこを経由してAPIを呼び出す必要がありましたが、それも必要ありません。
しかし、クライアント側でAPIキーを直接埋め込む必要があるため、扱いには注意が必要です。
信頼できるユーザに限定して公開するケースもしくは
ユーザ自身でAPIキーを設定するケースになるかと思います。

実装方法

実装は簡単で、APIのリクエストヘッダに

anthropic-dangerous-direct-browser-access: true`

を追加するだけです。
あとはx-api-keyヘッダにAPIキーを指定すればAPIを呼び出すことができます。

PGlite?

PGliteはElectricSQL社が開発した、PostgreSQLをWASMバイナリにコンパイルしたものです。
サイズも小さく(3Mくらい)、ブラウザやNode.jsなどで動作可能です。
また、多くのPostgreSQL拡張機能をサポートしており、
メモリ内DBとしても使用できますし、もちろん永続化も可能です。

import { PGlite } from "@electric-sql/pglite";
import { vector } from "@electric-sql/pglite/vector";

const db = await PGlite.create({ extensions: { vector } });

PGliteではpgvector拡張機能もデフォルトで使用可能になっており、
ここにあるように、vector型の使用や類似性検索がすぐに使えます。

Environment

MacBook Pro (14-inch, M3, 2023)
OS : MacOS 14.5
Node : v22.5.1
gh : v2.55.0 (2024-08-20)

Try

では、PGliteとClaude APIで、クライアントサイドだけで動作するRAGのデモをつくってみましょう。

今回は架空の会社の歴史をPGLiteに登録し、それをvector検索することで
社史とそのときの世間の出来事や背景もあわせて解説してくれるRAGアプリを実装してみます。

RAGアプリの動作フロー

デモアプリの流れは下記です。

アプリ起動時に PGlite データベースが初期化される
「データベースセットアップ/リセット」ボタンをクリックすると、サンプル社史データがPGLiteに登録される
AnthropicのAPIキーを入力し、質問を入力して「検索」ボタンをクリック
入力された質問に最も関連する会社の出来事をPGLiteから検索
検索結果と質問を使用して Claude API にリクエストを送信し、AI 生成の回答を取得
生成された回答が画面に表示される

アプリのリポジトリはこちらです。

セットアップ方法

↑のリポジトリからcloneして
下記のようにコマンドを実行すれば動きます。

# gh
% gh repo clone https://github.com/nakamura-shuta/pglite-react-app
% cd pglite-react-app
% npm i

# pglite.wasmをコピー
% mkdir -p public/@electric-sql/pglite/dist/
% cp ../node_modules/@electric-sql/pglite/dist/postgres.wasm public/@electric-sql/pglite/dist/

# Build & Run
% npm run build
% npm run dev

各コンポーネント

デモはシンプルなReactアプリです。
特徴的な部分をいくつか紹介。

社史データ登録

セットアップボタンを押すと、PGLiteに（テキストデータをvector化した）会社の歴史データを登録します。
※vector化処理はデモ用

//src/service/databaseOperation.js
export const insertCompanyHistory = async (db, history) => {
  await db.exec('DELETE FROM company_history');
  console.log('Existing data deleted');

  for (const item of history) {
    const embedding = getEmbedding(`${item.year}年: ${item.event}`);
    await db.query('INSERT INTO company_history (year, event, embedding) VALUES ($1, $2, $3)', 
                   [item.year, item.event, JSON.stringify(embedding)]);
    console.log(`Added event for year ${item.year}`);
  }
  console.log('Company history data added successfully');
};

Vector Search

検索ボタンを押すと、PGLiteに対して類似性に基づいた検索をします。

//src/service/databaseOperation.js
export const findRelevantEvents = async (db, query) => {
  const queryEmbedding = getEmbedding(query);
  const result = await db.query('SELECT year, event, embedding FROM company_history');

  return result.rows
    .map(row => ({
      ...row,
      distance: cosineSimilarity(queryEmbedding, JSON.parse(row.embedding))
    }))
    .sort((a, b) => a.distance - b.distance)
    .filter((event, index) => index < 3 || event.distance < 0.8);
};

Claude3 API

Claude3 APIを呼び出すコードです。
headerに必要情報を設定するだけで呼び出せます。

//src/services/aiOperations.js
export const getAIResponse = async (apiKey, query, eventDataList) => {
    const eventsText = eventDataList
      .map(event => `${event.year}年: ${event.event}`)
      .join('\n');

    const response = await fetch("https://api.anthropic.com/v1/messages", {
      method: "POST",
      headers: {
        "x-api-key": apiKey,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
        "anthropic-dangerous-direct-browser-access": "true",
      },
      body: JSON.stringify({
        model: "claude-3-sonnet-20240229",
        max_tokens: 1000,
        messages: [
          {
            role: "user",
            content: [<プロンプト>],
          },
        ],
      }),
    });

    const data = await response.json();
    console.log('AI response generated successfully');
    return data.content[0].text;
  };

実行結果

動かしてみます。

% npm run dev

> pglite-react-app@0.0.0 dev
> vite

  VITE v5.4.2  ready in 672 ms

  ➜  Local:   http://localhost:5173/
  ➜  Network: use --host to expose
  ➜  press h + enter to show help

データベースセットアップボタンを押し、
テスト用社史を登録してから検索します。

APIキーを入力し、下記の質問をします。
「2015年から2018年の出来事を教えてください」
しばらくすると、下記の回答が返ってきました。