langgraph-supervisorをDatabricks上で試す

導入

LangGraphのVersion 0.3が公開され、合わせて様々な事前ビルド済みエージェントライブラリが以下のページで紹介されています。

今回はその中のlanggraph-supervisorをDatabricks上で試します。

以前から気になっていたのですが、QwQ-32Bの実力を測るのに丁度よさそうだと思ったので、Quickstartを中心に実施してみます。

環境はDatabricks on AWS、クラスタはサーバレスです。

langgraph-supervisorとは

序段を邦訳して抜粋すると以下の通り。

LangGraphを使用して階層型マルチエージェントシステムを作成するためのPythonライブラリ。階層型システムは、中央のスーパーバイザーエージェントによって調整される特殊なエージェントのタイプのマルチエージェントアーキテクチャです。スーパーバイザーは、現在のコンテキストとタスクの要件に基づいて、どのエージェントを呼び出すかを決定し、すべての通信フローとタスクの委任を制御します。

特徴

• 🤖 複数の特殊エージェントを調整するスーパーバイザーエージェントを作成
• 🛠️ エージェント間の通信のためのツールベースのエージェント引き継ぎメカニズム
• 📝 会話制御のための柔軟なメッセージ履歴管理

このライブラリは、エージェントアプリケーションを構築するための強力なフレームワークであるLangGraphの上に構築されており、ストリーミング、短期および長期メモリ、人間を介在させるサポートをすぐに利用できます。

誤解を恐れず言えば、LangGraph上で構築した複数のエージェント(ReActエージェント)を状況に応じて使い分ける判断を行うエージェントを(階層的に)構築する仕組ですね。
ユーザはこのスーパーバイザエージェントとやりとりするだけで、裏側のエージェントの使い分けを気にする必要なく一連の処理を実行できることになります。また全体の短期/長期記憶の管理も担ってくれます。

余談。
このモジュールのコードも読んでみましたが、スーパーバイザエージェント自身も一種のReActエージェントとして作られています。また、エージェントへのタスク渡し/戻しをTool Callingで実装するなど、エージェント構築のアイディアが詰まったような作りでしたので、コード自体を一読してみるのをお薦めします。ボリュームもそこまで大きくはないため。

では、実際にQuickstartを動かしながら動作を確認してみます。

Quickstartを実行

以下の画像のように、数学関連の問題解決を行う"Math Agent"と検索を行う"Search Agent"、それらを使い分けて問題解決するスーパーバイザの三つのエージェントを使ってユーザの指示を実行するデモになります。

公式ではOpenAI社のGPT-4oを用いるようにしていましたが、今回は以下の記事で作成した、Databricks Model ServingにデプロイしたAribaba社のQwQ-32Bを利用します。

では、まずノートブックを作成して必要なパッケージをインストール。

%pip install langgraph-supervisor databricks-langchain
%pip install mlflow-skinny[databricks]
%pip install loguru

%restart_python

任意ですが、MLflow Tracingを有効化します。

import mlflow

mlflow.langchain.autolog()

また、思考経路を可視化するために、ChatDatabricksをオーバーライドしたクラスを用意。
エンドポイントから返ってきたcustom_outputsの中身も結果に含めています。

from typing import Any, Mapping
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from databricks_langchain import ChatDatabricks

class ChatDatabricksExt(ChatDatabricks):
    def _convert_response_to_chat_result(
        self, response: Mapping[str, Any]
    ) -> ChatResult:
        """Convert the response to a ChatResult."""
        resp = super()._convert_response_to_chat_result(response)

        if "custom_outputs" in response:
            resp.llm_output.update({"custom_outputs": response["custom_outputs"]})

        return resp

では、Quickstartを一部修正した内容を実行します。

from databricks_langchain import ChatDatabricks

from langgraph_supervisor import create_supervisor
from langgraph.prebuilt import create_react_agent

model_name = "sglang_qwen_qwq_32b_endpoint"
model = ChatDatabricksExt(model=model_name).bind(
    temperature=0.6, top_p=0.95, max_tokens=3000 # 公式のオススメ設定
)

# エージェントが実行するツールを関数として準備
def add(a: float, b: float) -> float:
    """Add two numbers."""
    return a + b


def multiply(a: float, b: float) -> float:
    """Multiply two numbers."""
    return a * b


def web_search(query: str) -> str:
    """Search the web for information."""
    return (
        "Here are the headcounts for each of the FAANG companies in 2024:\n"
        "1. **Facebook (Meta)**: 67,317 employees.\n"
        "2. **Apple**: 164,000 employees.\n"
        "3. **Amazon**: 1,551,000 employees.\n"
        "4. **Netflix**: 14,000 employees.\n"
        "5. **Google (Alphabet)**: 181,269 employees."
    )

# 各エージェントの作成
math_agent = create_react_agent(
    model=model,
    tools=[add, multiply],
    name="math_expert",
    prompt="You are a math expert. Always use one tool at a time.",
)

research_agent = create_react_agent(
    model=model,
    tools=[web_search],
    name="research_expert",
    prompt="You are a world class researcher with access to web search. Do not do any math.",
)

# スーパバイザの作成。内部的にはcreate_react_agentで作成されたLangGraphエージェントの一種。
workflow = create_supervisor(
    [research_agent, math_agent],
    model=model,
    prompt=(
        "You are a team supervisor managing a research expert and a math expert. "
        "For current events, use research_agent. "
        "For math problems, use math_agent."
    ),
)

# コンパイルして実行
app = workflow.compile()
result = app.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "what's the combined headcount of the FAANG companies in 2024?",
            }
        ]
    }
)

MLflow Tracingで結果を見てみます。

最終回答を拡大すると以下のようになります。

この回答は、コード内のweb_search関数が返す内容に加えて、各社のヘッドカウント総数が計算された内容になっていますね。

結果は正しそうなので、処理を簡単にMLflow Tracingの内容を基に追ってみます。

• ①エージェント実行後、最初にスーパーバイザで問い合わせ内容が処理され、問い合わせ内容からresearch_expertへハンドオフ(タスクの引継ぎ)が決定されます。transfer_to_researdch_expoertというツールが呼び出されることで、グラフの遷移が行われました。

• ②research_expertでweb_searchツールが実行され、FAANGの情報が取得されました。
  元情報にはヘッドカウント総数が記載されていないのですが、このタイミングで総数計算もされていますね。。。その上でtransfer_back_to_supervisorツールが実行され、スーパーバイザにタスクが戻ります。

• ③スーパーバイザで追加のエージェント処理を行うかどうかの判断がなされて、(十分な情報が得られているため)最終回答が実施されました。

これ、GPT-4oを使うと、FAANGのヘッドカウント総数を計算するためにMath Agentが呼ばれるのですが、今回のモデルだと最初から総数計算をしてしまっていたためここで処理終了となります。Reasoningモデルだからこれぐらいの計算はしてしまうのですね。。。

以上、Quickstartのウォークスルーでした。

まとめ

langgraph-supervisorをQwQ-32Bを使ってQuickstartをウォークスルーしてみました。
他にもいろいろ試しているのですが、複数エージェントに対するルーティングなど手軽に行えて面白いと思います。

個人的にマルチエージェントでなければならない/向いているユースケースがまだイメージ出来ていないところがあるのですが、実践を通じて見極めていく中で今回のようなライブラリがあるのは有難いと感じています。

Reasoningモデルとの組み合わせは、適正なハンドオフ先決定がされる、かつ理由が確認しやすいなど便利ではあるのですがレイテンシの大きさがボトルネックになりそう。Resoningモデルとそうでないモデルを的確に組み合わせるのが良さそうです。（多くの事に当てはまりますが）

マルチエージェントシステムに関するあれこれもキャッチアップしていろいろ試していきたいと思います。