Faiss-node + LangChainでEmbeddingしたベクトルデータをAmazon S3に保存・ロードする

FAISSを利用することで、OpenAIやClaudeなどでEmbeddingしたベクトルデータをローカルファイルとして保存・読み込みができるようになります。個人的にはベクトルDB版のSQLite的な手軽さがあるのが非常にありがたいOSSです。ただ人間は欲深いもので、「ファイルとして保存できる」となると今度は「Amazon S3に置きたいな」となってきました。

Amazon S3にFAISSのインデックスを配置するための試行錯誤を以下にまとめます。

現状のFAISS-nodeでは、ローカルにしかファイルを配置できない

そもそもライブラリにオプションがあればそれで解決・・・なのですが、コードを見るとそもそもローカルにファイルを配置することしか想定していない様子でした。fsを使っているので、Cloudflare Workers / Pages functionsで使うのも難しそうです。

  /**
   * Saves the current state of the FaissStore to a specified directory.
   * @param directory The directory to save the state to.
   * @returns A Promise that resolves when the state has been saved.
   */
  async save(directory: string) {
    const fs = await import("node:fs/promises");
    const path = await import("node:path");
    await fs.mkdir(directory, { recursive: true });
    await Promise.all([
      this.index.write(path.join(directory, "faiss.index")),
      await fs.writeFile(
        path.join(directory, "docstore.json"),
        JSON.stringify([
          Array.from(this.docstore._docs.entries()),
          this._mapping,
        ])
      ),
    ]);
  }

車輪の再開発で挙動を把握する

コードを読む限り、saveもloadもそれ自体はすごく難しい処理をしている様子もありませんでした。と、いうことは挙動を再現できればS3にインデックスを配置することもできそうです。

インデックスを保存する処理を自力で実装してみた

ということで、まずはsaveを使わずに同様の処理を実装してみました。ディレクトリ生成やライブラリのimport処理・パス指定などのオプションは考慮せず、シンプルに実装します。

    const vectorStore = await FaissStore.fromTexts(
      ["Hello world", "Bye bye", "hello nice world"],
      [{ id: 2 }, { id: 1 }, { id: 3 }],
      new OpenAIEmbeddings({
        openAIApiKey: process.env.OPENAI_API_KEY
      })
    );
    const faissIndex = vectorStore.index.toBuffer()

    fs.writeFileSync('./vector-store/manual/faiss.index', faissIndex);
    fs.writeFileSync('./vector-store/manual/docstore.json', JSON.stringify([
        Array.from(vectorStore.docstore._docs.entries()),
        vectorStore._mapping,
    ]));

docstore.jsonがすこし変わったスキーマになっている点に注意するくらいで、要はFaissStoreをJSONとBufferでそれぞれエクスポートしているだけです。fs.writeFileSyncをAWS S3のputObjectに変えれば動きそうですね。

インデックスを読み込む処理も再開発する

ファイルを配置できても、読み込めなければ使えません。ということでload側も再開発してみました。利用するライブラリはこの辺りです。

import { FaissStore } from "langchain/vectorstores/faiss";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { SynchronousInMemoryDocstore } from "langchain/stores/doc/in_memory"
import path from 'path';
import fs from 'fs';

そして実装側はこのような形になりました。2つのファイルを読み込み、FaissStoreクラスを作る処理をベタ書きで実装しています。

faissApp.get('load', async c => {
    const storeJSON = fs.readFileSync('./vector-store/manual/docstore.json', 'utf-8')
    const indexData = fs.readFileSync('./vector-store/manual/faiss.index')

    const [docstoreFiles, mapping] = JSON.parse(storeJSON)
    const docstore = new SynchronousInMemoryDocstore(new Map(docstoreFiles));

    const { default: {
        IndexFlatL2
    }} = await import('faiss-node')
    const index = IndexFlatL2.fromBuffer(indexData)
    const vectorStore = new FaissStore(
        new OpenAIEmbeddings({
            openAIApiKey: env.openai.apiKey
        }),
        {
            docstore,
            index,
            mapping
        }
    )
    const resultOne = await vectorStore.similaritySearch("hello world", 1);
    return c.json(resultOne)
});

こちらも「ファイルの中身が意図した形式で読めれば、fs以外からでも読み込みできる」実装であることが伺えます。

Amazon S3(AWS SDK v3)と接続する

動きがだいたいわかったので、いよいよS3とつなぎます。まずはSDKをいれましょう。

npm i @aws-sdk/client-s3

保存先をS3に変更する

先ほどのコードをS3のPutObjectに差し替えるだけです。オブジェクトストレージなので、ディレクトリ生成処理まわりを意識せずに済むのはよさそうですね。

    const vectorStore = await FaissStore.fromTexts(
      ["Hello world", "Bye bye", "hello nice world"],
      [{ id: 2 }, { id: 1 }, { id: 3 }],
      new OpenAIEmbeddings({
        openAIApiKey: env.openai.apiKey
      })
    );
    const faissIndex = vectorStore.index.toBuffer()
    await s3.send(
        new PutObjectCommand({
            Bucket: "S3_BUCKET_NAME",
            Key: "langchain/faiss/faiss.index",
            Body: faissIndex
        })
    )

    await s3.send(
        new PutObjectCommand({
            Bucket: "S3_BUCKET_NAME",
            Key: "langchain/faiss/docstore.json",
            Body: JSON.stringify([
                Array.from(vectorStore.docstore._docs.entries()),
                vectorStore._mapping,
            ])
        })
    )

S3に保存したデータを読み込みする

読み込み側も同様にfsをs3:GetObjectに差し替えればOKです。GetObjectで取得したデータを、docstore.jsonはtransformToString&JSON.parseで、faiss.indexはtransformToByteArray&Buffer.concat()で処理しています。

    const s3DocstoreData = await s3.send(
        new GetObjectCommand({
            Bucket: "S3_BUCKET_NAME",
            Key: "langchain/faiss/docstore.json",
        })
    )
    const docstoreData = await s3DocstoreData.Body?.transformToString()
    const [docstoreFiles, mapping] = JSON.parse(docstoreData || '')
    const docstore = new SynchronousInMemoryDocstore(new Map(docstoreFiles));
    const s3IndexData = await s3.send(
        new GetObjectCommand({
            Bucket: "S3_BUCKET_NAME",
            Key: "langchain/faiss/faiss.index",
        })
    )
    if (!s3IndexData.Body) {
        return c.json([])
    }
    const { default: {
        IndexFlatL2
    }} = await import('faiss-node')
    const indexData = await s3IndexData.Body?.transformToByteArray()
    const index = IndexFlatL2.fromBuffer(Buffer.concat([indexData]))
    const vectorStore = new FaissStore(
        new OpenAIEmbeddings({
            openAIApiKey: env.openai.apiKey
        }),
        {
            docstore,
            index,
            mapping
        }
    )

やってみて

faiss-nodeが各種プロパティをprivateなどでアクセスできないようにしているわけではなかったため、外からでもsaveとloadの挙動を再現できたのがかなり助かりました。ライブラリ自体がTypeScriptではない様子なので、そもそも設定できない様子だということもありますが・・・

本体側に取り込んでもらえると嬉しい気持ちもありますが、いろいろとベタ書きしている部分をオプション化したり、他のメソッドとの整合性を考えたりしないとですので、Pull Requestについては少し考えてから出すようにしたいと思います。

あと、理論上Cloudflare R2とかfirebase / GCPでもやれると思いますので、どなたかぜひ挑戦してみてください。

参考

• https://transang.me/modern-fetch-and-how-to-get-buffer-output-from-aws-sdk-v3-getobjectcommand/