LangchainとDatabricksで(私が)学ぶLCEL: 入力まわりの挙動

# 導入

LangchainでChainを定義する際に、[LCEL(LangChain Expression Language)](https://python.langchain.com/docs/expression_language/)という記法で書くことが推奨されてからだいぶ時間が経ったと思います。

実際、以前のLLMChain等を使う書き方と比べても柔軟性や記述の見通しがよく、Stream処理も書きやすいため、個人的には気に入っています。（RのTidyVerseと近いパイプ処理なところが馴染んでいるという側面もある）

ただ、結構なんとなく使っている部分もあって、```RunnablePassthrough```や```RunnableLambda```の挙動がイマイチ理解できていなかったり、把握していな機能もいろいろありそうだと感じていました。

気づいたら公式DocのLCEL HowTo部分がかなり充実していたので（前から？）、その内容を基に入力周りの挙動を確認しながら学んでみます。

https://python.langchain.com/docs/expression_language/how_to/

挙動確認はDatabricks on AWS上で実施しました。が、Databricks以外でも同様に実行できると思います。

# 単純なChainの確認

まず、シンプルに通常のChainをLCELで記述してみます。

必要なパッケージをインストール。

```python
%pip install -U langchain

dbutils.library.restartPython()
```

次に**プロンプトテンプレートだけ**で構成されるChainを作成。
通常Chainは、プロンプトテンプレート、LLM(Chat Model)、OutputParserで最低限構成しなければならないように見えますが、```Runnable```インターフェースを持つものであれば単体でChainを構成できます。

```python
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template("Tell me a short joke about {topic}")

chain = prompt
```

実行すると以下のような感じ。

```python
chain.invoke({"topic": "AI"})

# ChatPromptValue(messages=[HumanMessage(content='Tell me a short joke about AI')]) が出力される。
```

プロンプトテンプレートのパラメータtopicに値が設定されたChatPromptValueオブジェクトを取得できました。

では、これを拡張していきます。

# Chainへの入力を単一の文字列だけに変更する

chainの```invoke```メソッドで、辞書型のデータを渡していますが、いちいちキー名を付けて渡すのが面倒です。
value部分だけ渡して実行するとどうなるでしょうか。

```python
chain.invoke("AI")

# TypeError: Expected mapping type as input to ChatPromptTemplate. Received <class 'str'>. が発生
```

残念ながら、プロンプトテンプレートの処理で、キー名のマッピングができずエラーがでます。

下記リンク先の表にあるように、Chainを構成する要素は種類ごとにインプットの種類とアウトプットの種類が規定されています。
プロンプトテンプレートに対しては辞書型(Dictionary)をインプットとして与える必要があるのですが、```invoke```メソッドで文字列を渡すとそのままプロンプト側にも文字列がそのまま渡るため、エラーとなります。

https://python.langchain.com/docs/expression_language/interface

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/cb6b5247-b1cf-a054-83ec-1d815999b925.png)

この場合に使えるのが```RunnablePassthrough```です。

以下のようにすると、プロンプトに辞書型のデータを連携することができます。

```python
from langchain_core.runnables import RunnablePassthrough

chain = {"topic": RunnablePassthrough()} | prompt
```

実行してみましょう。

```python
chain.invoke("AI")

# ChatPromptValue(messages=[HumanMessage(content='Tell me a short joke about AI')])
# エラーなく実行できる
```

```RunnablePassthrough```はchainに入力された内容をそのまま得ることができます。
Chainの最初で辞書型のデータを作り、キーをtopic、値をRunnablePassthroughから取得することで辞書型に入力を変換しています。

これはRAGを実装するときに、Retrieverに検索文字列を与える際に有用です。
この場合のRAG Chainは以下のようなイメージになります。

```python
retrieval_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | model
    | StrOutputParser()
)

retrieval_chain.invoke("where did harrison work?")
```

Retrieverは、上の表にあるように単一の文字列しかインプットに指定できません。
そのため、Chainの呼び出し時に単一文字列を指定し、Chainの中で辞書型に変換するような処理を書くことで後続のプロンプトに適切なパラメータを渡すことが出来ます。

# Chainへの入力を複数パラメータにする

単一の文字列入力は簡単ですが、複数のパラメータをChainに与える際は、辞書型のパラメータを```invoke```に与えることになります。
一番最初の例では単一のキーを持つ辞書型データを使いましたが、複数のキーを持つ辞書型データを渡すこともできます。

```python
prompt = ChatPromptTemplate.from_template("{a} is {b}.")
chain = prompt

chain.invoke({"a": "Databricks", "b": "Intelligence Platform"})

# ChatPromptValue(messages=[HumanMessage(content='Databricks is Intelligence Platform.')])
```

このように複数のキーがある辞書型データの場合、```itemgetter```を使うことで個別の要素を取得することができます。
下の例では、当初のインプットに含むキー名がaとbだったものをSとCに詰め替えています。

```python
from operator import itemgetter

prompt = ChatPromptTemplate.from_template("{S} is {C}.")

# 入力データからa, bの要素を取り出して別のキーとして詰め替えた辞書型データを作成
chain = {
    "S": itemgetter("a"),
    "C": itemgetter("b"),
} | prompt

chain.invoke({"a": "Databricks", "b": "Intelligence Platform"})

# ChatPromptValue(messages=[HumanMessage(content='Databricks is Intelligence Platform.')])
```

ラムダ式など関数を指定することで、複数のパラメータを加工して新たなパラメータを作ることができます。
下の例では、与えられた二つのパラメータをisで繋ぎ、SVCというキーを持つ新たな辞書型データを作ってプロンプトに渡します。

```python
prompt = ChatPromptTemplate.from_template("{SVC}")

# ラムダ式（関数）を指定して、新たなパラメータを作る
chain = {
    "SVC": lambda x: x["a"] + " is " + x["b"]
} | prompt

chain.invoke({"a": "Databricks", "b": "Intelligence Platform"})
# ChatPromptValue(messages=[HumanMessage(content='Databricks is Intelligence Platform')])
```

他にも、RunnablePassthroughの```assign```メソッドを使うことで、入力された辞書型データを拡張することができます。
下の例は、上の例と同様にa, bという二つのパラメータをisで繋いた文字列をSVCというキーで追加していますが、もともと存在しているa, bのパラメータを失うことなく後続のプロンプトに渡しています。

```python
prompt = ChatPromptTemplate.from_template("{a} and {b} -> {SVC}")

# キーを変えて詰め替え
chain = RunnablePassthrough().assign(SVC=lambda x: x["a"] + " is " + x["b"] ) | prompt

chain.invoke({"a": "Databricks", "b": "Intelligence Platform"})

# ChatPromptValue(messages=[HumanMessage(content='Databricks and Intelligence Platform -> Databricks is Intelligence Platform')])
```

この例では、入力```{"a": "Databricks", "b": "Intelligence Platform"}```に対して```RunnablePassthrough().assign```を実行した結果、promptに渡されるのは```{"a": "Databricks", "b": "Intelligence Platform", "SVC": "Databricks is Intelligence Platform"}```となります。

<br>

上記で```itemgetter```で入力要素を個別に取り出す方法と、ラムダ式（関数）で要素を加工・追加しました。
では、これらを組み合わせた処理を作る際はどうしたらよいでしょうか。
例えば、```itemgetter```で取り出した文字や数字要素に対して、なんらかの加工処理を実施し、プロンプトに渡すようなケースです。

例えば、以下のようなやり方ができます。

```python
from langchain_core.runnables import RunnableLambda

prompt = ChatPromptTemplate.from_template("{sentence}")

def dummy_func(a):
    return a + "!!!!"

# itemgetterで取り出した要素をdummy_funcで加工する
chain = {
    # "sentence": itemgetter("a") | dummy_func # これだとエラー    
    "sentence": itemgetter("a") | RunnableLambda(dummy_func)
} | prompt

chain.invoke({"a": "Databricks", "b": "Intelligence Platform"})
# ChatPromptValue(messages=[HumanMessage(content='Databricks!!!!')])
```

```itemgetter```で取り出した要素を、LCELと同様のChainにつなぐ形で書くことが出来ます。
ただ、関数をそのままをパイプの後処理に指定することはできず、```RunnableLambda```でラップする必要があります。

これらを組み合わせると、複数パラメータを入力する際のRAG Chainを以下のように書くことが出来ます。

```python
# ダミーのRetrieverの役割を果たす関数を定義
def search_retriever(query: str):
    return f"Dummy context for {query}"


prompt = ChatPromptTemplate.from_template("Context:{context} and Question:{question}.")

chain = (
    # keywordを小文字化した要素を追加
    RunnablePassthrough().assign(lower_keyword=lambda x: x["keyword"].lower())
    # lower_keywordをRetrieverに渡し、questionはそのままquestionとして詰め替える
    | {
        "context": itemgetter("lower_keyword") | RunnableLambda(search_retriever),
        "question": itemgetter("question"),
    }
    | prompt
)

chain.invoke({"keyword": "Databricks", "question": "What is lakehouse?"})
# ChatPromptValue(messages=[HumanMessage(content='Context:Dummy context for databricks and Question:What is lakehouse?.')])
```

処理を組み合わせることで、入力が単一文字列でも、複数キーの辞書型データであっても、柔軟にRAG用のChainを定義できますね。

# おまけ：RunnableParallelを使わなくていいの？

上記の最後の例は、以下のように書くこともできます。

```python
from operator import itemgetter
from langchain_core.runnables import (
    RunnableParallel,
    RunnablePassthrough,
    RunnableLambda,
)

# ダミーのRetrieverの役割を果たす関数
def search_retriever(query: str):
    return f"Dummy context for {query}"


prompt = ChatPromptTemplate.from_template("Context:{context} and Question:{question}.")

# キーを変えて詰め替え
chain = (
    # keywordを小文字化した要素を追加
    RunnablePassthrough().assign(lower_keyword=lambda x: x["keyword"].lower())
    # **RunnableParallel**でラップ
    | RunnableParallel(
        {
            "context": itemgetter("lower_keyword") | RunnableLambda(search_retriever),
            "question": itemgetter("question"),
        }
    )
    | prompt
)

chain.invoke({"keyword": "Databricks", "question": "What is lakehouse?"})
# ChatPromptValue(messages=[HumanMessage(content='Context:Dummy context for databricks and Question:What is lakehouse?.')])
```

違いは、```RunnableParallel```というクラスで辞書型データをラップしているところです。

```RunnableParallel```はラップした処理を並列に実行するという定義をChainに加えます。
上記例で言えば、```itemgetter("lower_keyword") | RunnableLambda(search_retriever)```と```itemgetter("question")```の処理を並列に実行するように設定されます。
これによって、Retrieverの処理に時間がかかるときでも他の処理を並列で実行することができ、処理効率が上がります。

ということは、なるべく```RunnableParallel```でラップする方がよい、ということになるのですが、以下の公式DocのTips部分を見ると、自動で辞書型データはRunnableParallelでラップされるため、明示的に記述する必要がなさそうです。

https://python.langchain.com/docs/expression_language/how_to/map

というわけで、今回は```RunnableParallel```を明示的に使用はしませんでした。
とはいえ、LCELで大事な概念の一つだと思いますので、理解しておくことは重要だと思います。

# まとめ

LCELでChainを記述する際の、入力周りについて挙動を確認してみました。
Langchainは当初Chainの種類が増えすぎてて訳が分からなくなりがちだったのですが、LCELで記述できるようになってかなりスッキリしたと感じています。
```astream_log```メソッドを使うことで詳細な処理の流れも追いやすいです。

今回は非常に基本的かつベタな部分ではありますが、改めてDocを確認しながら実践するのは大事だなと思いました（小並感）

あとLangchainのドキュメント、結構整備され直している（気がする）。
core部分だけでもきちんと読み直していこう。