langchainとDatabricksで(私が)学ぶAgent : LangGraphでAgent処理を組む

導入

以下の記事で、ローカルLLMを使ってAgent(ReAct)の実践をしてみました。

今回は、LangGraphを使って、この処理を書き直してみます。

LangGraphとは？

こちらの公式Blogを参照ください。

npaka大先生が邦訳されたものを上げておられます。

一部引用すると、

「LangGraph」は、LLMでステートフルな「マルチアクターアプリケーション」を構築するためのライブラリです。「LCEL」(LangChain Expression Language) を拡張して、複数チェーン (またはアクター) を複数ステップにわたって循環的に協調動作させることができます。

「LangChain」の大きな価値の1つに、カスタムチェーンを簡単に作成できることがあります。このための機能として「LCEL」を提供してきましたが、サイクルを簡単に導入する方法がありませんでした。「LangGraph」によって、LLMアプリケーションにサイクルを簡単に導入できるようになりました。

LCELはサイクル(ループ)を記述するには向いていませんが、LangGraphを活用することでAgentに必要なサイクルを記述・導入できるようになります。

というわけで、やってみます。
今回はほぼ以下の公式Exampleのウォークスルーになります。

実践・検証はDatabricks on AWS上で実施しました。

Step1. パッケージのインストール

使うパッケージをインストール。
今回もExLlama V2を使って推論します。
また、langgraphを追加でインストールしています。

%pip install -U transformers accelerate "exllamav2>=0.0.11" langchain langchainhub duckduckgo-search

dbutils.library.restartPython()

Step2. モデルのロード

LLMのモデルを読み込みます。
今回もこちらで作成した、langchainのカスタムチャットモデルを使って読み込みました。

使っているモデルはOpenChat 3.5の0106版をGPTQで量子化した以下のモデルです。
事前にダウンロードしておいたものを使います。

from exllamav2_chat import ChatExllamaV2Model

model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--openchat-3.5-0106-GPTQ"

chat_model = ChatExllamaV2Model.from_model_dir(
    model_path,
    cache_max_seq_len=8192,
    system_message_template="GPT4 Correct User: {}<|end_of_turn|>GPT4 Correct Assistant: OK.<|end_of_turn|>",    
    human_message_template="GPT4 Correct User: {}<|end_of_turn|>GPT4 Correct Assistant: ",
    ai_message_template="{}",
    temperature=0.0001,
    top_p=0.0001,
    max_new_tokens=1024,
    repetition_penalty = 1.15,
    low_memory=True,
    cache_8bit=True,
)

# TheBloke/openchat-3.5-0106-GPTQはtokenizerのEOS設定に不具合があるので修正
chat_model.exllama_tokenizer.eos_token_id = 32000

Step3. ツールの準備

LLM Agentの中で呼び出して実行するツールを定義します。

前回同様、DuckDuckGoを使ってWeb検索するツールを使うことにします。

from langchain_community.tools.ddg_search.tool import DuckDuckGoSearchResults, DuckDuckGoSearchRun

search = DuckDuckGoSearchRun(max_results=1)

# Agentに利用するツールのリスト
tools = [search]

Step4. Agent用Chainの作成

LangChain HubからReActエージェント用のプロンプトテンプレートを取得し、create_json_chat_agent関数を使ってAgentのChainを作成します。

from langchain import hub
from langchain.agents import create_json_chat_agent

prompt = hub.pull("hwchase17/react-chat-json")

# Construct the ReAct agent
agent_runnable = create_json_chat_agent(chat_model, tools, prompt)

ここまでは前回とほぼ同じです。

Step5. エージェント状態の定義

ここからがLangGraphを使った処理となります。

まず、Agent内の状態を保持するためのクラスを定義します。
Agent処理中の内部状態はこのクラスのインスタンスに保持されます。
このクラスは後ほどLangGraphの内部で利用します。

状態として保持する情報は以下の通り。

• input: ユーザからの入力内容
• chat_history: Agent実行前の会話履歴
• intermediate_steps: Agent処理中の中間実行内容・結果
• agent_outcome: agentの応答結果、AgentActionかAgentFinishのインスタンスが格納される。応答結果がAgentFinishだと処理終了させるべきで、それ以外だとtoolを実行させるようなハンドリングとなる。

from typing import TypedDict, Annotated, List, Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.messages import BaseMessage
import operator

class AgentState(TypedDict):
   # ユーザーが入力した文章
   input: str
   # 会話の履歴メッセージのリスト
   chat_history: list[BaseMessage]
   # エージェント呼び出しの結果
   # Noneが有効なタイプとして設定する必要がある
   agent_outcome: Union[AgentAction, AgentFinish, None]
   # アクションとそれに対応する観測のリスト
   # ここで `operator.add` を注釈付きで使うことで、 `state` への操作が既存の値に追加されることを示している
   intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]

Step6. ノード用関数の定義

Agentに実行させる処理を定義します。
これらの関数がGraphのノードと対応します。

ここで定義するノード用の関数は以下の通り。

• run_agent: Step4で定義したagentのchainを実行する処理
• execute_tools: 指定されたツールを実行する処理

should_continue関数は直前の応答結果の内容を確認して、Agentの処理を終了するべきか、継続するべきかのハンドリングに利用します。

from langchain_core.agents import AgentFinish
from langgraph.prebuilt.tool_executor import ToolExecutor

# ツールを実行するのに便利なヘルパークラスです。
# エージェントアクションを取り込み、そのツールを呼び出して結果を返します
tool_executor = ToolExecutor(tools)

#  エージェントを定義する
# `agent` で追加されたキーであることに注意してください
# エージェントアウトカムを取得します。
def run_agent(data):
    agent_outcome = agent_runnable.invoke(data)
    return {"agent_outcome": agent_outcome}

# ツールを実行する関数を定義する
def execute_tools(data):
    # 一番最新の agent_outcome を取得する - これは agent で追加されたキーです
    agent_action = data['agent_outcome']
    output = tool_executor.invoke(agent_action)
    return {"intermediate_steps": [(agent_action, str(output))]}

# 条件分岐でどちらのエッジを進むかを決定するために使用されるロジックを定義する
def should_continue(data):
    # エージェントの出力が AgentFinish の場合は、 'end' 文字列を返します
    # これは、グラフをセットアップしてフローを定義する際に使用されます
    if isinstance(data['agent_outcome'], AgentFinish):
        return "end"
    # それ以外の場合は、AgentAction が返される
    # ここで 'continue' 文字列を返します
    # これは、フローを定義する際に使用されます
    else:
        return "continue"

通常はこれだけでいいのですが、今回利用しているモデルはあまりToolを使ってくれない感じなので、最初に必ずToolを使うようにするためのノード用関数first_agentを追加定義します。

from langchain_core.agents import AgentActionMessageLog

# ツール 'duckduckgo_search' を実行するための代理アクションを作成する。
# 入力のキー 'input' をそのまま渡す。
def first_agent(inputs):
    action = AgentActionMessageLog(
        # 強制的にこのツールを呼び出す
        tool="duckduckgo_search",
        # このツールに 'input' キーをそのまま渡す
        # 本当は入力を検索用キーワードに変換する処理を入れた方がよい
        tool_input=inputs["input"],
        log="",
        message_log=[],
    )
    return {"agent_outcome": action}

Step7. グラフの定義

Agent処理のグラフを定義します。
グラフは大まかにノードとエッジで構成されます。
このあたりはWikipediaのグラフ理論を見るとイメージが付きやすいと思います。

from langgraph.graph import END, StateGraph

# 新しいグラフを定義
workflow = StateGraph(AgentState)

# まず、三つのノードを定義する
# 第一引数がノード名、第二引数が対応する処理の関数
workflow.add_node("agent", run_agent)
workflow.add_node("action", execute_tools)
workflow.add_node("first_agent", first_agent)

# エントリポイントを「エージェント」に設定
# これにより、最初にこのノードが呼び出される 
workflow.set_entry_point("first_agent")

# 次にエッジの追加。
# まず条件分岐するエッジを追加する
workflow.add_conditional_edges(
    # 開始ノード。この処理ではagentノードから次のどこに進むかを定義している
    "agent",
    # 分岐先を決定する関数
    should_continue,
    # should_continueの結果に基づく終端ノード。
    # 一致するものに基づいて、そのノードが呼び出されます。
    # ENDは、グラフが終了する特別なノード    
    {
        # should_continueの結果がcontinueの場合
        "continue": "action",
        # should_continueの結果がendの場合
        "end": END,
    },
)

# actionノード -> agentノード　への通常のエッジ
workflow.add_edge("action", "agent")
# first_agent -> actionノード　への通常のエッジ
workflow.add_edge("first_agent", "action")

# 最後に、グラフをコンパイル
# これにより、LangChain Runnableにコンパイルされる
app = workflow.compile()

これでAgentを実行する準備ができました。

Step8. 実行

これまで定義したAgent用のグラフを実行します。
コンパイル後はLangChain Runnableのオブジェクトなので、LCELのChainと同様に利用することができます。

というわけで実行してみましょう。

inputs = {"input": "what is the weather in Tokyo", "chat_history": []}
for s in app.stream(inputs):
    print(list(s.values())[0])
    print("----")

出力

{'agent_outcome': AgentActionMessageLog(tool='duckduckgo_search', tool_input='what is the weather in Tokyo', log='', message_log=[])}
----
{'intermediate_steps': [(AgentActionMessageLog(tool='duckduckgo_search', tool_input='what is the weather in Tokyo', log='', message_log=[]), "Tokyo (Japan) weather - Met Office Today 11° 3° Sunny. Sunrise: 06:50 Sunset: 16:52 M UV Thu 18 Jan 13° 6° Fri 19 Jan 13° 4° Sat 20 Jan 9° 4° Sun 21 Jan 12° 6° Mon 22 Jan 14° 6° Tue 23 Jan Today... Today's weather forecast for Tokyo by the hour. Why Treksplorer? Founded in 2011 by Ryan O'Rourke, Treksplorer provides travel recommendations and advice to millions of readers every year. Our content is rooted in our writers' firsthand experiences, in-depth research, and/or collaborations with other experts and locals. Read more about our editorial policy. Weather in Tokyo: An overview 简体中文 ไทย Tiếng Việt 한국어 Indonesian Special Features Area Things to Do Mountains Sports & Outdoors Castles Temples Onsens & Spas Museums & Art Galleries Japan at night Flowers & Trees Shrines Tourist Spots & Attractions Theme Parks Nature & Scenery Parks & Gardens Food & Drinks Donburi & rice Dietary restrictions & preferences Meat Fish & Seafood MeteoState: Be prepared with the most accurate and detailed weather forecast for Tokyo with high temperature, low temperature, precipitation, dew point, humidity, wind speed and direction, atmospheric pressure, relative humidity, including current conditions, setting and rising of the Sun and Moon, Moon phase, geomagnetic forecast, state of the sea, wave height, assistance in the selection of ...")]}
----
{'agent_outcome': AgentFinish(return_values={'output': 'The weather in Tokyo today is sunny with a high of 11°C and a low of 3°C.'}, log='\n\n{\n"action": "Final Answer",\n"action_input": "The weather in Tokyo today is sunny with a high of 11°C and a low of 3°C."\n}')}
----
{'input': 'what is the weather in Tokyo', 'chat_history': [], 'agent_outcome': AgentFinish(return_values={'output': 'The weather in Tokyo today is sunny with a high of 11°C and a low of 3°C.'}, log='\n\n{\n"action": "Final Answer",\n"action_input": "The weather in Tokyo today is sunny with a high of 11°C and a low of 3°C."\n}'), 'intermediate_steps': [(AgentActionMessageLog(tool='duckduckgo_search', tool_input='what is the weather in Tokyo', log='', message_log=[]), "Tokyo (Japan) weather - Met Office Today 11° 3° Sunny. Sunrise: 06:50 Sunset: 16:52 M UV Thu 18 Jan 13° 6° Fri 19 Jan 13° 4° Sat 20 Jan 9° 4° Sun 21 Jan 12° 6° Mon 22 Jan 14° 6° Tue 23 Jan Today... Today's weather forecast for Tokyo by the hour. Why Treksplorer? Founded in 2011 by Ryan O'Rourke, Treksplorer provides travel recommendations and advice to millions of readers every year. Our content is rooted in our writers' firsthand experiences, in-depth research, and/or collaborations with other experts and locals. Read more about our editorial policy. Weather in Tokyo: An overview 简体中文 ไทย Tiếng Việt 한국어 Indonesian Special Features Area Things to Do Mountains Sports & Outdoors Castles Temples Onsens & Spas Museums & Art Galleries Japan at night Flowers & Trees Shrines Tourist Spots & Attractions Theme Parks Nature & Scenery Parks & Gardens Food & Drinks Donburi & rice Dietary restrictions & preferences Meat Fish & Seafood MeteoState: Be prepared with the most accurate and detailed weather forecast for Tokyo with high temperature, low temperature, precipitation, dew point, humidity, wind speed and direction, atmospheric pressure, relative humidity, including current conditions, setting and rising of the Sun and Moon, Moon phase, geomagnetic forecast, state of the sea, wave height, assistance in the selection of ...")]}
----

Agentの処理(グラフの各ノード)結果が順に出力されていきます。
最終行内のagent_outcomeの内容が得られた最終結果となります。

特に実行ログが不要な場合は、invokeで最後の結果のみ取得できます。
別のクエリを実行し、最終結果だけ出力してみます。

result = app.invoke({"input": "what is Databricks", "chat_history": []})

print(result["agent_outcome"].return_values["output"])

出力

Databricks is a cloud-based platform for managing and analyzing large datasets using the Apache Spark open-source big data processing engine. It offers a unified workspace for data scientists, engineers, and business analysts to collaborate, develop, and deploy data-driven applications.

まとめ

LangGraphを使ったAgentの実装でした。

今回の実行例では処理が循環しない（一方向に進んでそのまま終了している）ので、LangGraphを使う意味が薄いのですが、LLMの応答→ツールの実行→その結果を使ったLLMの応答→別のツールの実行...というAgentを組む場合には、非常に強力なフレームワークになるかと思います。

機能が既にあるのかどうか把握できていないのですが、グラフを可視化する機能もあるといいなあと感じています。

マルチエージェント用のサンプル等も用意されているようなので、このあたりもう少し掘り下げていきたいと思います。