RunnableK8sをbatch()とstream()対応したけど使い道がわからない

これは、[Kubernetes Advent Calendar 2024](https://qiita.com/advent-calendar/2024/kubernetes)の10日目の記事で、以前投稿した「[LangChainのRunnableをKubernetesで実行するRunnableを書いてみた](https://qiita.com/kaitoy/items/6da99a974da5a4588617)」という記事の続き。

# RunnableK8sとは
[RunnableK8s](https://github.com/kaitoy/runnable-k8s/tree/master)は、LangChainのRunnalbeをラップすると、その部分をKubernetesで実行できるRunnable。
以下のような感じで、ラップしたRunnableとそのインプットをシリアライズして、Podの標準入力に流し込んで渡して、Pod内で実行し、結果をPodの標準出力経由で返す、という風に動く。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/110166/2fb7a7f9-b28f-2b78-c6b4-aca26d25ed54.png)

[前回の記事](https://qiita.com/kaitoy/items/6da99a974da5a4588617)の時点ではRunnableの最も基本的なメソッドである`invoke()`だけを実装していた。
今回の記事では、`batch()`と`stream()`を実装して動かしてみる。

# batch()とstream()とは
`invoke()`が、ひとつのインプットを受け取り、Runnableを同期的に一回一括で実行して一つのアウトプットを生成するのに対して、`batch()`と`stream()`は以下のような感じ。

* `batch()`: 複数のインプットを受け取り、(並列で)Runnableを同期的にインプットの数だけまとめて実行して、複数のアウトプットを生成する。
* `stream()`: 一つのインプットを受け取り、Runnableを同期的に1回実行しつつ、アウトプットをストリームとして生成する。もう少し正確に言うと、`stream()`を実行すると、アウトプットの断片(chunk)を順次生成するイテレータを返す。

# batch()の実装
`batch()`は自前のコードを書く必要がなかった。
Runnableクラスに実装されている`batch()`の[デフォルト実装](https://github.com/langchain-ai/langchain/blob/v0.1.12/libs/core/langchain_core/runnables/base.py#L497)が、インプットの数だけスレッドを立ち上げて`invoke()`を呼ぶというもので、それで十分なので継承したものをそのまま採用。

この`batch()`を実行するサンプルコードは以下。

```python:sample_batch.py
import sys

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from runnables.runnable_k8s import RunnableK8s


prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        "あなたはポケモンのエキスパートです。ユーザの質問に答えてください。",
    ),
    (
        "human",
        "質問: {input}",
    )
])
llm = ChatOpenAI(model="gpt-4o-mini")

chain = prompt | RunnableK8s(bound=llm) | StrOutputParser()

outputs = chain.batch([
  {'input': sys.argv[1]},
  {'input': sys.argv[2]},
  {'input': sys.argv[3]},
])

for o in outputs:
    print(o + '\n')
```

このコードは、コマンドラインから受け取った3つのポケモンに関する質問について、プロンプトを生成し、OpenAIに投げて、回答をテキストに変換して標準出力に吐く、というのを`batch()`でまとめて実行する。
OpenAIに投げる部分をRunnableK8sでラップしているので、そこだけKubernetesで実行される。

これを実行した様子は以下。

![batch.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/110166/87ff93d1-da59-058a-8963-07a39e59406e.gif)

RunnableK8sが3つのpodを同時に作って、質問を並列で処理できている。
回答は全部間違ってるけど。

# stream()の実装
`stream()`は、Runnableクラスの[デフォルト実装](https://github.com/langchain-ai/langchain/blob/v0.1.12/libs/core/langchain_core/runnables/base.py#L567)が単に`invoke()`を呼び出すというもので、ストリーム処理になってないので、[自前で実装](https://github.com/kaitoy/runnable-k8s/blob/1e45ef2476c440fa2693e9ef924448c902e0cdc3/app/runnables/runnable_k8s.py#L305-L316)した。

実装の中身は、冒頭の図のような処理をする`invoke()`と大きく変わらないが、podのなかでRunnableの`stream()`を実行し、そのchunkを順次標準出力に吐き、RunnableK8s側のイテレータで順次返すようにした。

この`stream()`を実行するサンプルコードは以下。

```python:sample_stream.py
import sys

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from runnables.runnable_k8s import RunnableK8s


prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        "あなたはポケモンのエキスパートです。ユーザの質問に答えてください。",
    ),
    (
        "human",
        "質問: {input}",
    )
])
llm = ChatOpenAI(model="gpt-4o-mini")

chain = prompt | RunnableK8s(bound=llm) | StrOutputParser()

for output in chain.stream({'input': sys.argv[1]}):
    print(f'{output}/', end='', flush=True)
print()
```

このコードは、コマンドラインから受け取った1つのポケモンに関する質問について、プロンプトを生成し、OpenAIに投げて、回答をテキストに変換して標準出力に吐く、というのを`stream()`でストリーム実行する。

RunnableK8sでラップしたOpenAIのLLMからは回答がchunkで分割されて順次くるので、それを標準出力に吐くときに、chunkの間に`/`を挟んで切れ目をわかりやすくしている。

これを実行した様子は以下。

![stream.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/110166/1c633703-8429-eedc-f9b7-e8b9b8fc2a7d.gif)

RunnableK8sでストリーム処理できた。
回答もちょっとくどいけど合ってる。