LangChainのRunnableBranchで条件分岐した処理をする

やりたいこと

前のChainで得られた結果によって呼ぶChainを変えたい

RunnableBranch

RunnableBranchを使うと

Runnable that selects which branch to run based on a condition.
The Runnable is initialized with a list of (condition, Runnable) pairs and a default branch.

conditionとRunnableのペアをリストとDefault branchを設定すればいいということがわかります。

from langchain_core.runnables import RunnableBranch

branch = RunnableBranch(
    (lambda x: isinstance(x, str), lambda x: x.upper()),
    (lambda x: isinstance(x, int), lambda x: x + 1),
    (lambda x: isinstance(x, float), lambda x: x * 2),
    lambda x: "goodbye",
)

branch.invoke("hello") # "HELLO"
branch.invoke(None) # "goodbye"

例として載っているのは、入力の値の型によって、処理を変え、デフォルトは"goodbye"と返すものです。

結構わかりやすいですね。

Example: メッセージがあればサマリーし、なければ固定で no message found を返す

具体的に自分がやりたかったことは、Slack MessageのSummaryですが、Slack APIから取得した結果が空のときは、そもそもSummary するChainをよばずに固定の値を返して、最後のメッセージを変えたいと思っていました。

以下の例では、メッセージを取得する部分とサマリーする部分は擬似的にRunnableLambdaで表現しています。

from random import randint
from langchain_core.runnables import RunnableLambda, RunnableBranch


def example_runnable_branch():
    # branch
    extract_message_chain = RunnableLambda(lambda x: {"messages": ["message"] * randint(0, 2)})  # randomize message
    message_summary_chain = RunnableLambda(lambda x: f"this is the summary of {len(x['message'])} messages")  # emulate chain with LLM

    # chain
    chain_with_branch = extract_message_chain | RunnableBranch(
        (lambda x: bool(x["messages"]), message_summary_chain),
        lambda x: "no message found",  # fixed result when empty
    )
    print(chain_with_branch.batch([{}, {}, {}]))

• extract_message_chain でmessageの数が異なるmessageが生成される
• message_summary_chain でmessageの数を数えてサマリーしたことにしている
• RunnableBranch で、 x["messages"]が存在するときにはmessage_summary_chainを呼び、デフォルトでは"no message found"を返す

実行結果:

毎回ランダムなので結果は変わってしまいますが、no message foundやsummaryされたタイプの結果が返ってきていて想定通り処理できている事がわかりました。

['this is the summary of 1 messages', 'this is the summary of 1 messages', 'no message found']