LangChain(TypeScript)のRetrievalQAChainとOpenAI APIで、自前のドキュメントに関する質問に答えてくれるプログラムを作ってみた。

概要

OpenAIのChatGPTさん、めっちゃ頭良いけど、自分しか持って無くてインターネットに公開していない情報については回答してくれないので、自分しか持ってないドキュメントの内容に関する質問に回答してくれるプログラムを作ってみたので、その備忘録。

GitHubリポジトリ

こちらのリポジトリにソースコードはアップロードしています。

github.com

使う技術

• Node.js（バージョン18.14.2）

• TypeScript

• LangChain（TypeScript/JavaScript 版）

• OpenAI API（APIキーを発行します）

ダミーファイルを作成

そんなすぐに、良い感じのドキュメントを用意できないので、BingのAIさんに、いくつか社内用語を出力してもらって、それをテキストファイルとして保存しました。

テキストファイルはすべて documentsディレクトリに保存しています。

ディレクトリ構成

こんな感じ

create-vector フォルダ

documentsフォルダ内のテキストをvectorストアに保存するプログラム

call-query フォルダ

保存済のvectorストアを読み込み、質問の回答を得るプログラム

documents フォルダ

テキストファイルを保存しておくフォルダ（ここでは架空の企業の社内用語と説明のテキストファイル）

database フォルダ

ベクターストアの内容を保存しておくフォルダ

今回はベクターストアとして、ローカルに保存できるHNSWLibを利用します。

実行の流れ

ドキュメントの内容をベクターストアに保存するプログラムを実行する ↓ 保存済のvectorストアを読み込み、質問の回答を得るプログラムを実行する

プログラムのコード

create-vector-store/index.ts

documentsフォルダ内のドキュメントの内容をベクターストアに保存するプログラム

RecursiveCharacterTextSplitterで、ドキュメントの内容を100token以下になるまで分割しVectorStoreに保存しています。

import path from 'path';
import dotenv from 'dotenv';
import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
import { DirectoryLoader } from 'langchain/document_loaders/fs/directory';
import { TextLoader } from "langchain/document_loaders/fs/text";
import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
import { HNSWLib } from "langchain/vectorstores/hnswlib";


// .envファイルから環境変数を読み込み
dotenv.config();


const main = async function () {
  console.log("start");

  // 自前で準備したドキュメントを保存しているディレクトリの絶対パス
  const documentPath = path.join(__dirname, '../documents/');

  const directoryLoader = new DirectoryLoader(documentPath, {
    '.txt': (path) => {
      return new TextLoader(path)
    }
  });
  const documents = await directoryLoader.load();

  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 100,
    chunkOverlap: 10,
  });

  const texts = await splitter.splitDocuments(documents);
  const embeddings = new OpenAIEmbeddings({
    openAIApiKey: process.env.OPENAI_API_KEY
  });
  const vectorStore = await HNSWLib.fromDocuments(texts, embeddings);

  const databasePath = path.join(__dirname, '../database/');
  await vectorStore.save(databasePath);

  console.log("end");
};

main();

call-query/index.ts

保存済のvectorストアを読み込み、質問の回答を得るプログラム

ここで、保存済のベクターストアから読み取り、RetrievalQAChainを使って自前のドキュメントの内容に対する質問に関する質問の回答を得ます。

import path from 'path';
import dotenv from 'dotenv';
import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
import { OpenAI } from "langchain/llms/openai";
import { HNSWLib } from "langchain/vectorstores/hnswlib";
import { RetrievalQAChain } from "langchain/chains";


// .envファイルから環境変数を読み込み
dotenv.config();


const main = async function () {
  // 実行時引数が指定されているか確認
  if (process.argv.length < 3) {
    return;
  }

  const databasePath = path.join(__dirname, '../database/');
  const embeddings = new OpenAIEmbeddings();

  // 保存済のベクターストアから読み込み
  const vectorStore = await HNSWLib.load(databasePath, embeddings);

  const model = new OpenAI({
    openAIApiKey: process.env.OPENAI_API_KEY
  });
  const chain = RetrievalQAChain.fromLLM(model, vectorStore.asRetriever());

  // 実行時引数から質問内容を取得
  const query = process.argv[2];

  const { text } = await chain.call({ 
    query
  });

  // 質問内容と回答内容をコンソールに表示
  console.log(`Q: ${ query }\nA: ${ text }`);

};

main();

実行手順

ベクターストアの作成

$ cd <プロジェクトディレクトリ>/create-vector-store/
$ npm run start

質問をする

$ cd <プロジェクトディレクトリ>/create-vector-store/
$ npm run start --query=<質問>

実行結果例

ちゃんとdocumentsフォルダに保存しているドキュメントの内容に沿った回答をもらえました。

また、試しにdocumentsフォルダ内のドキュメントには記載の無い質問をすると、ちゃんと "分からない" 旨の回答がかえってきました。素直だな。

まとめ

とりあえず、自前で準備したドキュメントに関する回答を得ることができた。

今回はDirectoryLoaderの中で、テキストファイルに対してのみ読み取りを行っているが、これをPDFやWord、PowerPointに対応させれば、より汎用的に使えるようになるかもしれないなと思いました。