FakeLLMを使って、LangChainのChainを無料・高速にテストする方法

LangChainでは、複数のLLM API呼び出しやデータ処理をサポートする仕組みとしてChainが提供されています。Chainを利用することで、「1回目の生成結果を利用して、新しいテキスト生成を行う」や「入力文に応じて、参照するデータソースを変更する」などの動きを実装することができます。下の例では、「人名と言語を指定すると、『その人の出身地に関する文章』を、指定した言語で生成する処理」を、2回のLLM API呼び出しで実現させています。

const prompt1 = PromptTemplate.fromTemplate(
  `What is the city {person} is from? Only respond with the name of the city.`
);
const prompt2 = PromptTemplate.fromTemplate(
  `What country is the city {city} in? Respond in {language}.`
);

const model = new ChatAnthropic({});

const chain = prompt1.pipe(model).pipe(new StringOutputParser());

const combinedChain = RunnableSequence.from([
  {
    city: chain,
    language: (input) => input.language,
  },
  prompt2,
  model,
  new StringOutputParser(),
]);

const result = await combinedChain.invoke({
  person: "Obama",
  language: "German",
});

Chainの検証・テストには、FakeLLMを使おう

Chainを利用することで、LLMのレスポンスの整形などの手間を省くことができます。ただしLLMのAPIを呼び出す関係上、ローカルでの試行錯誤や高頻度のCI / CDパイプライン上でのテストの実施は、意図しない出費につながる可能性があります。

そんな時に利用できるのが、LangChainが提供する「FakeLLM」です。このモデルは、その名の通り偽のLLMを提供します。テキスト生成を行うモデルを利用しないため、このモデルを利用する際は、「どんなテキストを返すか」を事前に定義します。

const llm = new FakeListLLM({
  responses: ["I'll callback later.", "You 'console' them!"],
});

配列で定義したものは、前から順番に返却されます。配列の長さより多い回数の呼び出しがあった場合は、配列の先頭に戻る形でループします。

FakeLLMを使うと便利なユースケース（一例）

FakeLLMを使うか使わないかの判断は、「何をテストしたいか」で行います。例えばOutputParserでのレスポンスの整形や、Chainに渡すinput情報の整形や処理、Retrieverへのクエリだけテストしたい場合などには、LLMで都度テキストを生成する必要がないため、FakeLLMが便利です。一方で「テキスト生成」に関する処理やChainをテストしたい場合は、LangSmithなどを利用する方が良いかなと思います。

ChainをInjectableにして、テストを書きやすくする

LLMを本物とFakeLLM双方使えるようにする方法の例として、Chainを抽象化してモデルをInjectするサンプルを紹介します。

class ExampleChain {
  constructor(
    private readonly model: BaseChatModel
  ) {}
  async invoke(question: string) {
    const prompt = PromptTemplate.fromTemplate(
      `Given the user question below.
      
      Question: {question}
      Answer:`
    )
  
    const chain1 = RunnableSequence.from([
      {
        question: new RunnablePassthrough()
      },
      prompt,
      this.model,
      new StringOutputParser()
    ])
    return chain1.invoke(question)
  }
}

このChainを例えばCloudflare Workers AIで利用する場合は、次のように実装します。

  const chat =  new ChatCloudflareWorkersAI({
    model: "@cf/meta/llama-2-7b-chat-int8",
    cloudflareAccountId: c.env.CLOUDFLARE_ACCOUNT_ID,
    cloudflareApiToken: c.env.CLOUDFLARE_API_TOKEN,
  });

  const chain = new ExampleChain(chat)
  const r = await chain.invoke("Who are you?")

そしてこのChainに対してユニットテストを書きたい場合は、FakeLLMを渡してやりましょう。

describe('ExampleChain', () => {
  it('should return a response from FakeListChatModel', async () => {
    const fakeChat = new FakeListChatModel({
      responses: ["I'll callback later."],
    });

    // Mocking ExampleChain
    const chain = new ExampleChain(fakeChat);

    // Invoke the chain
    const response = await chain.invoke("Who are you?");

    // Assertion
    expect(response).toEqual("I'll callback later.");
  });
});

「テキスト生成」以外の処理は、FakeLLMを検討しよう

FakeLLMを使うことで、LangChainが持つ「テキスト生成以外」の機能を安く手軽に試すことができます。LLMを呼び出さずに処理を走らせることができるのは、試行錯誤の多いケースやCI / CD上でテストを走らせたい場合にもありがたいのではないでしょうか。

一方で「モデルを差し替えしやすくする実装方法」が開発時に要求される点には注意が必要です。環境変数で出し分けるよりは、モデルを明示的にコンストラクタなどで注入する形にする方が、Flaskyなコード・テストになりにくくてよいかと思います。

関連記事

• https://js.langchain.com/docs/integrations/llms/fake
• https://js.langchain.com/docs/integrations/chat/fake