Amazon Aurora Serverless v2にpgvectorをインストールしてOpenAIのembeddingデータを格納する

はじめに

OpenAIに始まり、AWSでも2023/9/29にAmazon BedrockがGAしたりと、2023年はLLM関連のニュースが非常に多くなっている。その中でも、Vector DBと組み合わせてRAGを用いることでハルシネーションの回避をすることができるようになるため、非常に重要な技術になってきている。
AWSでも、2023/7/13よりAurora(PostgreSQL互換)がpgvectorをサポートして、よりVector DBが身近なものになった。

今回は、このpgvector拡張のインストールと使い方を確認していく。

Amazon Aurora(PostgreSQL互換)を準備する

まずは、Amazon Auroraクラスタを起動しよう。
Amazon Auroraクラスタの起動は過去の記事を参考にしていただきたい。
なお、pgvectorのサポートしているPostgreSQLのエンジンバージョンは上記のWhat's Newに記載の通り、以下のバージョン以降でないと動作しないため、注意が必要だ。過去の記事の内容からバージョンは変更していただきたい。

pgvector 拡張機能は、AWS GovCloud (米国) リージョンを含む AWS リージョンにおいて、Aurora PostgreSQL 15.3、14.8、13.11、12.15 以降で利用できます。

起動後、接続して以下のSQLを実行する。

create_extention.sql

create extension vector;

さらに、ベクトルデータを格納するためのテーブルと索引を作成する。

create_table.sql

create table documents (
  id bigserial primary key,
  content text,
  embedding vector(1536)
);

create index on documents using ivfflat (embedding vector_l2_ops)
with
  (lists = 100);

索引の作成については正直よく分かっていない。
このページによると、距離演算の方法によって作成する索引がことなるらしい。今回、ユークリッド距離の演算で索引を作成してみる(コサイン演算も試してみたが安定しなかった)。

これで、準備は完了だ。

Open AIのembeddingデータを格納する。

以下のTypescriptでデータを格納をしてみよう。
今回は、過去の想い出をcontentsに格納して、それをopenai.createEmbedding()で、text-embedding-ada-002のモデルを使ってベクトルデータ化している。

import { Client } from 'pg';
import pgvector from 'pgvector/pg';
import { Configuration, OpenAIApi } from 'openai';
import config from 'config';

const contents: Array<string> = [
  '夏や年末になると実家に遊びに行く。実家には、婆ちゃんと爺ちゃんとペットの犬が住んでいて、いつも大歓迎してくれる。',
  'みんなでよく近所の公園に散歩に行っていた。大きな公園で、遊具場ではたくさんの子ども達が遊んでいる公園だ。',
  '小学校には給食室があって、給食センターからの配送ではなくそこで給食を作っていた。できたての給食はとても美味しかった。',
  '高校の頃は、毎日、友達とダべったり、バンドの練習室に行ったり、ゲームセンターに行ったりしていた。',
];

// Configuration for OpenAIApi
const configuration = new Configuration({ apiKey: config.get('credentials.openaiApiKey') });
const openai = new OpenAIApi(configuration);

export async function createEmbedding (): Promise<void> {
  try {
    const client = new Client({
      host: config.get('applicationSettings.auroraUrl'),
      port: config.get('applicationSettings.auroraPort'),
      database: 'VECTOR',
      user: config.get('applicationSettings.auroraUsername'),
      password: config.get('applicationSettings.auroraPassword'),
    });
    await client.connect();
    await pgvector.registerType(client);

    for (const content of contents) {
      const input = content.replace(/\n/g, ' ');

      const response = await openai.createEmbedding({
        model: 'text-embedding-ada-002',
        input,
      });

      const [{ embedding }] = response.data.data;

      await client.query('INSERT INTO documents (content, embedding) VALUES ($1, $2)', [
        input,
        pgvector.toSql(embedding),
      ]);
    }

    await client.end();
  } catch (error: any) {
    console.error(error);

    throw new Error('createEmbedding() failed.');
  }
}

(async () => {
  await createEmbedding();
})();

なお、await client.end();はしておかないと、メイン関数が終了しないので注意。

ベクトルデータ作成の際、改行を渡すことができないのでその数行前にreplace()を行っている。

格納した情報から検索を行ってみる

import { Client } from 'pg';
import pgvector from 'pgvector/pg';
import { Configuration, OpenAIApi } from 'openai';
import config from 'config';

// Configuration for OpenAIApi
const configuration = new Configuration({ apiKey: config.get('credentials.openaiApiKey') });
const openai = new OpenAIApi(configuration);

export async function dataSelect (queryContent: string): Promise<string> {
  try {
    const client = new Client({
      host: config.get('applicationSettings.auroraUrl'),
      port: config.get('applicationSettings.auroraPort'),
      database: 'VECTOR',
      user: config.get('applicationSettings.auroraUsername'),
      password: config.get('applicationSettings.auroraPassword'),
    });
    await client.connect();
    await pgvector.registerType(client);

    const input = queryContent.replace(/\n/g, ' ');

    const response = await openai.createEmbedding({
      model: 'text-embedding-ada-002',
      input,
    });

    const [{ embedding }] = response.data.data;

    const queryResult = await client.query('SELECT * FROM documents ORDER BY embedding <=> $1 LIMIT 1', [pgvector.toSql(embedding)]);

    await client.end();

    return queryResult.rows[0].content;
  } catch (error: any) {
    console.error(error);

    throw new Error('dataSelect() failed.');
  }
}

(async () => {
  const result = await dataSelect(process.argv[2]);
  console.log(`result: ${result}`);
})();

いざ、使ってみよう！

package.jsonには以下を定義しておく。

{
  (中略)
  "scripts": {
    (中略)
    "insert": "ts-node --files -r tsconfig-paths/register ./src/insert.ts",
    "select": "ts-node --files -r tsconfig-paths/register ./src/select.ts"
  }
  (中略)
}

必要なモジュールは、npm installで入れておこう。

で、npm run insertすると、データベースには以下の4つのレコードが格納されるはずだ。

VECTOR=> select id, content from documents;
 id |                                                       content                                                        
----+----------------------------------------------------------------------------------------------------------------------
  1 | 夏や年末になると実家に遊びに行く。実家には、婆ちゃんと爺ちゃんとペットの犬が住んでいて、いつも大歓迎してくれる。
  2 | みんなでよく近所の公園に散歩に行っていた。大きな公園で、遊具場ではたくさんの子ども達が遊んでいる公園だ。
  3 | 小学校には給食室があって、給食センターからの配送ではなくそこで給食を作っていた。できたての給食はとても美味しかった。
  4 | 高校の頃は、毎日、友達とダべったり、バンドの練習室に行ったり、ゲームセンターに行ったりしていた。

この後、以下のように検索をしてみる。

$ npm run select 実家
result: 夏や年末になると実家に遊びに行く。実家の家には、婆ちゃんと爺ちゃんとペットの犬が住んでいて、いつも大歓迎してくれる。

$ npm run select 近所の公園
result: みんなでよく近所の公園に散歩に行っていた。大きな公園で、遊具場ではたくさんの子ども達が遊んでいる公園だ。

$ npm run select 給食センター
result: 小学校には給食室があって、給食センターからの配送ではなくそこで給食を作っていた。できたての給食はとても美味しかった。

$ npm run select 高校の頃
result: 高校の頃は、毎日、友達とダべったり、バンドの練習室に行ったり、ゲームセンターに行ったりしていた。

こんな感じで一部のワードから類似したものを持ってきてくれるようだ。
例えば、ワードを変えて以下のようにしても良い感じに近いものを持ってきてくれる

$ npm run select バンドとゲームセンター
result: 高校の頃は、毎日、友達とダべったり、バンドの練習室に行ったり、ゲームセンターに行ったりしていた。

これで、OpenAIで発生する揺らぎを良い感じに吸収してくれるということだろう。
OpenAIのFunction CallingのFunctionで、今回作った検索の処理を動かしてあげれば、カスタマイズした外部記憶をAIに与えることができるようになるということだ。