💡

LangChainのソースコードリーディングをしてみたわかったこと

2023/09/04に公開

はじめに

個人的に、LangChainを使ったアプリケーション開発に興味を持っており、その挙動を理解するために公式ドキュメントだけでなくソースコードを読むことがありました。
この記事では、LangChainのソースコードリーディングをして得られた知見について紹介します。

LangChainとは？

LangChainはLLMを使ったアプリケーションを開発するときに使えるOSSのライブラリです。
Open AIはもちろん、Google PaLM、Hugging Face等の他のLLMのラッパーとして使用することもできます。

ソースコードリーディングした環境

LangChainというとPython版が主流ですが、自分はJS/TSの方が得意なので、JS/TS版のソースコードリーディングをしていきます。
バージョンは0.0.14を見ていきます。

学んだこと

複数のLLMが同一のインターフェースで使えるように設計されている

「LangChainとは？」の節でも触れましたが、複数のLLMが使用できるように設計されています。

例えば、人とLLMが対話するときに使用するAPIがOpenAIとGoogle PaLMには存在しまが、レスポンスのフォーマットが異なります。

それぞれ、抽象基底クラス（BaseChatModel）で抽象メソッドとして定義されている、 _generate メソッドの実装で差分を吸収しています。

https

  abstract _generate(
    messages: BaseMessage[],
    options: this["ParsedCallOptions"],
    runManager?: CallbackManagerForLLMRun
  ): Promise<ChatResult>;

[https://github.com/hwchase17/langchainjs/blob/0.0.140/langchain/src/chat_models/base.ts#L321-L325]

興味がある人は読んでみてください。

Adoptorパターンの手本のようなコードになっているなーと感じました。

いい感じのプロンプトでLLMのAPIを叩いてくれる

OpenAIはtokenの上限があるため、大量のドキュメントを要約させるためには工夫が必要です。

これに関して、LangChainではRefine、Map reduceといった方法が提供されており、再帰的に実行する、並列的に実行し最後に集約するといった方法で大量のドキュメントの要約を可能にしています。

再帰的に操作するといった処理自体はすぐに実装できそうですが、LLMの場合、プロンプトも大事になってきます。
確認してみたところ、LangChainではRefine（再帰的に処理する方法）のプロンプトは次のように定義されていました。

langchain/src/chains/summarization/refine_prompts.ts

const refinePromptTemplate = `Your job is to produce a final summary
We have provided an existing summary up to a certain point: "{existing_answer}"
We have the opportunity to refine the existing summary
(only if needed) with some more context below.
------------
"{text}"
------------

Given the new context, refine the original summary
If the context isn't useful, return the original summary.

REFINED SUMMARY:`;

[https://github.com/hwchase17/langchainjs/blob/0.0.140/langchain/src/chains/summarization/refine_prompts.ts]

操作だけでなく、プロンプトも含めて用意されており、loadSummarizationChain関数を呼び出すだけでそれがよしなに実行されるというのはすごいですね。
これはもうLangChain使うしかないなーと思いました。

人の実装の誤りを吸収してくれる仕組みが入ってる

OpenAI APIには「Chat Completions API」と「Completions API」という二つのAPIが存在します。
前者は後発に出たAPIになっており、後者の上位互換的存在だと思ってください（詳しくはこちら）。

後者のAPIは最新のGPT-4には存在しないAPIなのですが、なんと、LangChainで定義されている後者のAPIを操作するクラス（OpenAI）をGPT-4を指定してインスタンスを生成することができます。
見てみると次のようになっており、前者のAPIを操作するクラス（OpenAIChat）のインスタンスが生成されていることがわかります。

langchain/src/llms/openai.ts#L147-L154

    if (
      fields?.modelName?.startsWith("gpt-3.5-turbo") ||
      fields?.modelName?.startsWith("gpt-4") ||
      fields?.modelName?.startsWith("gpt-4-32k")
    ) {
      // eslint-disable-next-line no-constructor-return, @typescript-eslint/no-explicit-any
      return new OpenAIChat(fields, configuration) as any as OpenAI;
    }

[https://github.com/hwchase17/langchainjs/blob/0.0.140/langchain/src/llms/openai.ts#L147-L154]

本来開発者が気を付けないといけないところですが、「Completions API」から「Chat Completions API」への移行がしやすくなるように、設計されているのだと思います。

おわりに

LangChainを使うことで、素のAPIを使うよりも、アプリケーションに組み込みやすいように色々仕込まれていることを感じられたかと思います。
他にもChainやAgentといった便利なAPI、サードパーティのSaaS、ライブラリとのインテグレーション、サードパーティーライブラリの導入などなど、アプリケーションに組み込む上でしたいことを準備してくれています。

難しいことをしているように思っていましたが、インタフェース設計がしっかりしているので、コードは案外読みやすいというのが意外でした。
そして、LLMをどう扱ったらいいかわかったり、複数のAPIを同様に扱うための設計の参考としてなど、単純に勉強になりました。

公式ドキュメントでは拾いきれない細かい挙動については、ソースコードリーディングするのが手っ取り早いと感じました。