⛓️

RunnableLambdaの使い方と入力制約

2024/06/14に公開

LCELではRunnableクラスを継承した何かを、パイプで繋いでいくことでChainを表現する。
通常の関数はRunnableLambdaでラップすることでRunnableにできる。

環境

  • langchain 0.2.3
  • langchain-core 0.2.5

RunnableLambdaの基本的な使い方

まず、RunnableLambdaの基本的な使い方。例えば、整数を受け取って1を加える簡単な関数をラップする場合、以下のように実装する。

from langchain_core.runnables import RunnableLambda

def add_one(x: int) -> int:
    return x + 1

runnable_add_one = RunnableLambda(add_one)

result = runnable_add_one.invoke(3)
print(result)  # 出力: 4

ちなみに、以下のようにすると複数回Runnableを回せる。

result = runnable_add_one.map().invoke([1, 2, 3])
print(result)  # 出力: [2, 3, 4]

入力として1つの引数しか取れない制約

Runnableクラス全般において、基本的に入力として1つの引数しか取れないという制約がある。これは、Runnableオブジェクトがパイプラインやシーケンス内で連結して使われることを想定しているため。この設計により、複雑なデータ処理フローを構築する際の一貫性とシンプルさが保たれるのだが、この仕様に気づくまでコードと睨めっこした。

それでも、複数の引数を取る関数をRunnableLambdaしたい場合は多々ある。いくつか手段を考えてみた。

タプルを使用する

この方法では、lambda関数を使ってタプルを展開し、元の関数に渡す。

from langchain_core.runnables import RunnableLambda

def add_numbers(x: int, y: int) -> int:
    return x + y

runnable_add_numbers = RunnableLambda(lambda inputs: add_numbers(*inputs))

result = runnable_add_numbers.invoke((3, 4)) 
print(result) # 出力: 7

辞書を使用する

辞書も使える。PromptTemplateに複数の値を渡すコード例でよく目にする。

from langchain_core.runnables import RunnableLambda

def add_numbers(data: dict[str, int]) -> int:
    return data["x"] + data["y"]

runnable_add_numbers = RunnableLambda(add_numbers)

result = runnable_add_numbers.invoke({"x": 3, "y": 4}) 
print(result) # 出力: 7

タプルと辞書は、どちらもinvoke側からは中身の型が分からないのが難点。

状態を持たせる

ちなみに、クラスメソッドもRunnableにできる。クラスインスタンスで状態を持たせることができるので、いろいろできそう。

from langchain_core.runnables import RunnableLambda

class Adder:
    def __init__(self):
        pass

    def add(self, data: dict[str, int]) -> int:
        return data["x"] + data["y"]

adder = Adder()
runnable_adder = RunnableLambda(adder.add)

result = runnable_adder.invoke({"x": 3, "y": 4}) 
print(result) # 出力: 7

まとめ

RunnableLambdaで任意のPython関数をラップすることができれば、LCELの自由度が大きく向上する。シーケンシャルなフロー設計の場合、LLMではないプロダクトでも使えるかもしれない。
ちょっと試してみたい。

Discussion