LangChain.jsで作成済みのChainをToolとして登録する方法

この記事は「LangChain Advent Calendar 2023」24日目の記事です。

ToolsもChainも両方使いたい

Toolsを使うことで、例えばLangChainがサポートしていない任意のAPIをよびだして結果を生成するみたいな使い方ができます。しかしToolsを使うには、ChainではなくAgentExecutorを使う必要があります。

  const executor = await initializeAgentExecutorWithOptions(tools, model, {
    agentType: "zero-shot-react-description",
    verbose: true
  });

ChainをCustom Toolsとして登録する

これがベストプラクティスなのかはまだ自信がありません。ただ、ドキュメントなどをみる限りでは、Tools側に寄せる形で対応するが良さそうです。

まずはChainを用意しましょう。

    const model = new OpenAI({
        openAIApiKey: process.env.openaiApiKey,
        temperature: 0,
    })
    const reviewTemplate = `あなたは入力された文章から、検索キーワードを抽出する作業を行っています。入力されたメッセージから、キーワードをスペース区切りのテキストで出力してください。
出力するテキストには、「教えて」などの動詞は含めないでください。
input: {query}

{format_instructions}

`;
    const outputParser = StructuredOutputParser.fromNamesAndDescriptions({
      output: "Suggested search keyword",
    })
    const reviewPromptTemplate = new PromptTemplate({
          template: reviewTemplate,
          inputVariables: ["query"],
          partialVariables: {
              format_instructions: outputParser.getFormatInstructions()
          },
          outputParser: outputParser,
    });
    const queryChain = new LLMChain({
      llm: model,
      prompt: reviewPromptTemplate,
    });

その後、ChainをDynamicToolとして登録します。

  const tools = [
    new DynamicTool({
        name: "KeywordPickup",
        description: "入力されたテキストから検索キーワードを抽出します。キーワードは2つまで指定できます。",
        func: async (input) => queryChain.run(input),
    })
  ];

  const executor = await initializeAgentExecutorWithOptions(tools, model, {
    agentType: "zero-shot-react-description",
    verbose: true
  });

これでChainもToolsも一緒に動かすことができます。

ログで動作を見る

意図した通りに動いているかを、verbose: trueで出力したログから見てみましょう。

[chain/end] [1:chain:AgentExecutor > 5:chain:LLMChain] [654ms] Exiting Chain run with output: {
  "text": " これは何かの情報を抽出するためのツールですね。\nAction: KeywordPickup\nAction Input: Amazon Bedrock"
}
[agent/action] [1:chain:AgentExecutor] Agent selected action: {
  "tool": "KeywordPickup",
  "toolInput": "Amazon Bedrock",
  "log": " これは何かの情報を抽出するためのツールですね。\nAction: KeywordPickup\nAction Input: Amazon Bedrock"
}

みている限り、Chainについても意図した通りのToolとして認識してくれている様子です。

[llm/end] [1:chain:AgentExecutor > 2:chain:LLMChain > 3:llm:OpenAI] [755ms] Exiting LLM run with output: {
  "generations": [
    [
      {
        "text": " Amazon Bedrockとは何かを知りたいということですね。\nAction: KeywordPickup\nAction Input: Amazon Bedrock",
        "generationInfo": {
          "finishReason": "stop",
          "logprobs": null
        }
      }
    ]
  ],
  "llmOutput": {
    "tokenUsage": {
      "completionTokens": 32,
      "promptTokens": 178,
      "totalTokens": 210
    }
  }
}
[chain/end] [1:chain:AgentExecutor > 2:chain:LLMChain] [756ms] Exiting Chain run with output: {
  "text": " Amazon Bedrockとは何かを知りたいということですね。\nAction: KeywordPickup\nAction Input: Amazon Bedrock"
}
[agent/action] [1:chain:AgentExecutor] Agent selected action: {
  "tool": "KeywordPickup",
  "toolInput": "Amazon Bedrock",
  "log": " Amazon Bedrockとは何かを知りたいということですね。\nAction: KeywordPickup\nAction Input: Amazon Bedrock"
}

書き言葉（自然文）として入力されたテキストから、無事検索キーワードのピックアップに成功しました。

感想

「ユースケースは少なそうだけれども、知らないと困りそうかも？」と思ったのがやってみての感想です。基本的にはLangChainが用意してくれているRetriever / Loaderを使う方向で頑張りたいなとは思いますが、PoCなどで少しだけ独自の動きをつけたくなったケースなどに役立つかもしれません。