TypeScriptでLangChainを使ってみるその3 RAGパターン編

こんばんは。
最近Alexaの耳が遠くなってる気がするんですが、気のせいでしょうか。

前回は検索編でしたが、今回は応用としてRAGパターン編です。
shironeko.hateblo.jp

引き続きコードはGitHubにあるので良かったらご覧ください。
github.com

データの収集と登録

参照するデータがないとRAGにならないので、とりあえずデータを集めて保存することにします。
今回はサンプルとしてWikipediaの日本の記念日一覧 - Wikipediaというページをデータソースにしてみます。

汎用性もへったくれもないDocumentLoaderですが、まぁサンプルデータ集めるだけなんでね…。

class AnniversaryLoader extends BaseDocumentLoader {
    async load(): Promise<Document[]> {

        const wikipediaQuery = new WikipediaQueryRun({
            baseUrl: "https://ja.wikipedia.org/w/api.php"
        });

        const result = await wikipediaQuery.content("日本の記念日一覧");
        return [new Document({ pageContent: result })];
    }
}

そしてこのままだと1つのでかいドキュメントになってしまうのでいい感じに分割します。分割はTextSplitterですよね。
いやーなんとも雑なコードです。ページのレイアウトが変わったら一発でアウトですが、まぁいいでしょう。

class AnniversaryTextSplitter extends TextSplitter {
    splitText(text: string): Promise<string[]> {

        const lines = text.split("\n");

        let title = "";
        let day = "";
        let anniversaries: string[] = [];

        for (const line of lines) {
            if (line.match(/^== (.+) ==$/)) {
                title = line.split(" ")[1].trim();
            } else if (title.match(/^\d+月$/) && line.match(/^\d+日.+/)) {
                day = line.split(" ")[0];
                anniversaries.push(...line.split(" ")[2].split("、").map((a: string) => `${title}${day} : ${a}`));
            }
        }

        return Promise.resolve(anniversaries);
    }
}

で、こんな感じでChromaDBにベクトル化して登録します。

const embeddings = new OpenAIEmbeddings();
const vectorStore = new Chroma(embeddings, {
    url: process.env.CHROMADB_URL,
    collectionName: "anniversary"
});

const loader = new AnniversaryLoader();
const docs = await loader.loadAndSplit(new AnniversaryTextSplitter());

await vectorStore.addDocuments(docs);

これで記念日の情報がベクトル検索できるようになりました。
ちなみにですが、「今日の東京の天気」みたいなリアルタイムの情報は基本的にはVectorStoreに保存したりせずに直接APIとかを叩くのが良いかと思います。
リクエスト回数が多い場合などにはキャッシュとして保存するのはもちろんありですが。

さくっとRAGパターン

ここまでで構成する要素は一通り触ってきてるのであとは組み合わせるだけです。
今回はちょっとテスタブルを意識して依存性は引数で渡すようにしてみたので、Retriever・LLM・Promptを取得する関数を定義しておきます。

function getRetriever(): BaseRetriever {

    const embeddings = new OpenAIEmbeddings();
    const chroma = new Chroma(embeddings, {
        url: process.env.CHROMADB_URL,
        collectionName: "anniversary"
    });

    return chroma.asRetriever();
}

function getLLM(): BaseChatModel {
    return new ChatGoogleGenerativeAI();
}

function getChatTemplate(): ChatPromptTemplate {
    return ChatPromptTemplate.fromMessages([
        SystemMessagePromptTemplate.fromTemplate(`コンテキストに沿ってユーザの入力に回答してください。複数ある場合はすべて回答してください。
----------
{context}`),
        HumanMessagePromptTemplate.fromTemplate("{input}")
    ]);
}

で、コンソールから入力を受け取りたいのでこうします。連続したい場合は無限ループで回したりなどお好みで。

const reader = createInterface({ 
    input: stdin, 
    output: stdout
});

const question = await reader.question("記念日について質問をどうぞ：");
reader.close();

入力待ちになりますね。

> langchainsample@1.0.0 rag
> ts-node rag_execute.ts

記念日について質問をどうぞ：

順番に

まずは普通にRetrieverで検索して、LLMChainで回答させてみましょう。
シグネチャがないとややこしいかもしれないので、一旦関数定義ごと。入力の他にRetriever・LLM・Promptを引数で渡しています。

async function rag1(question: string, retriever: BaseRetriever, llm: BaseChatModel, chatTemplate: ChatPromptTemplate) {

    const docs = await retriever.getRelevantDocuments(question);
    const context = docs.map(d => d.pageContent).join("\n");

    const llmChain = new LLMChain({ prompt: chatTemplate, llm: llm });
    const result = await llmChain.invoke({ context, input: question });
    console.log(result.text);
}

ちゃんと答えてくれている気がします。でも元のページだと「税理士記念日」っていうのもあるんだけどな…。

記念日について質問をどうぞ：2/23はなんの日ですか？
* ふろしきの日
* 富士山の日

RetrievalQAChain

で、まぁこういうパターンってよく使うよねってことで実はRetrievalQAChainというChainが元々用意されています。
inputKeyやoutputKeyは時々出てくるんですが、パーツを組み合わせるときにこのキーの値を使うよみたいな宣言です。
複数のデータを渡した際にその内のどれを使うかなどを指定します。RetrievalQAChainだとデフォは「query」なのでテンプレートに合わせて「input」にしておきます。

async function rag2(question: string, retriever: BaseRetriever, llm: BaseChatModel) {

    const chain = RetrievalQAChain.fromLLM(llm, retriever, { inputKey: "input" });
    const result = await chain.invoke({ input: question });
    console.log(result.text);
}

あってはいるんですが、2/22は「猫の日」「食器洗い乾燥機の日」「忍者の日」「乃木坂46の日」「竹島の日」「行政書士記念日」らしいので淡白ですね。
実はこれ引数を見るとわかるのですが、テンプレートを渡さずにLangChainのテンプレートに任せています。

記念日について質問をどうぞ：2/22の記念日を教えて
2/22の記念日は「忍者の日」です。

llmのverboseをtrueにするとプロンプトが確認できるので見てみると、内容的には似たようなことが書いてありますが、すべて出力しろって言うニュアンスはないですね。
（ところで2/22の記念日他にもあるのに2/28のエッセイ記念日が混じってきてるのはなんでなの…）

(略)
        "kwargs": {
          "content": "Use the following pieces of context to answer the users question. \nIf you don't know the answer, just say that you don't know, don't try to make up an answer.\n----------------\n2月22日 : 竹島の日\n\n2月22日 : 猫の日\n\n2月28日 : エッセイ記念日\n\n2月22日 : 忍者の日",
          "additional_kwargs": {}
        }
(略)

テンプレートありRetrievalQAChain

テンプレートを利用してRetrievalQAChainを使うにはこんな感じです。

async function rag3(question: string, retriever: BaseRetriever, llm: BaseChatModel, chatTemplate: ChatPromptTemplate) {

    const chain = RetrievalQAChain.fromLLM(llm, retriever, { inputKey: "input", prompt: chatTemplate });
    const result = await chain.invoke({ input: question });
    console.log(result.text);
}

ついでにRetriverが取ってくるドキュメントの件数が5件だと足りなそうなので、10件にしておきます。asRetrieverの引数に件数を指定するだけです。

function getRetriever(): BaseRetriever {

    const embeddings = new OpenAIEmbeddings();
    const chroma = new Chroma(embeddings, {
        url: process.env.CHROMADB_URL,
        collectionName: "anniversary"
    });

    return chroma.asRetriever(10);
}

うーん…良くなってはいるけど全部は教えてくれないですね…。この辺はプロンプトの調整次第というところでしょうか。

記念日について質問をどうぞ：2/22の記念日を教えて
2/22の記念日は以下の通りです。

・竹島の日
・猫の日
・忍者の日
・乃木坂46の日

StuffDocumentsChainとRetrievalChain

LangChainのドキュメントを見ているとRetrievalQAChainはレガシーっぽいです。（なんならConversationalRetrievalQAChainはあるけどRetrievalQAChainは書いてすらない）
js.langchain.com

createStuffDocumentsChainとcreateRetrievalChainを使うのがいい感じなんでしょうか。

async function rag4(question: string, retriever: BaseRetriever, llm: BaseChatModel, chatTemplate: ChatPromptTemplate) {

    const documentChain = await createStuffDocumentsChain({ prompt: chatTemplate, llm: llm });
    const retrievalChain = await createRetrievalChain({combineDocsChain: documentChain, retriever: retriever});

    const result = await retrievalChain.invoke({ input: question });
    console.log(result.answer);
}

慣れればパズル

慣れてくるとコードは割と短いんですが、パラメタとかに何を設定するのかがよくわからなくて結局LangChainの中を毎回見てて時間がかかる…。
まぁいい勉強にはなるんですけどね。

いったん目的としていたRAGパターンを試すことができました。
うーん、次はAgentsに手を出すか…LCEL Chainsの書き方ももうちょっと整理したい気もする…。

白猫のメモ帳

C#とかJavaとかJavaScriptとかHTMLとか機械学習とか。