LangGraphのFunctional APIをDatabricksの機能と組み合わせて使ってみる

導入

LangGraphにFunctional APIが導入されました。(2025年2月時点でベータリリースのようです）

公式Blogでの紹介がこちら。

そして、既に詳しい日本語解説記事もQiitaに公開されていました。

Functional APIが何かというと、"通常のプログラミングパラダイムを用いてワークフローを定義し、通常の条件分岐制御やループ処理を行ってフローを制御できる"仕組です。(公式blogより)

LangGraphはその名前の通りグラフを構築することで（エージェント処理などの）ワークフローを定義していましたが、条件分岐を含むエッジやノードを独自の処理で構築する必要がありました。
これは単純なグラフでは良いのですが、複雑な条件分岐・ループを含んだグラフを作ろうと思うと結構大変だったというのが個人的意見だったのですが、そのあたりを通常のプログラム表現(Pythonであれば、Pythonのifやwhileループ表現）で記載できるようになります。

グラフを使わないのにLangGraphとはこれいかに。

とはいえ、これまでのグラフ(Graph API)を補完する位置づけのようで、適材適所で併用することが推奨されているように読み取れます。

詳細は公式ドキュメントや上記Qiita記事を見られるのが一番わかりやすいと思うので割愛し、Databricks上で以下のチュートリアルを一部改変して試してみます。

なお、改変部分は主にDatabricksの特徴に関する以下の検証を含んでいます。

• MLflowのTracingがFunctional API利用時にどうなるのか
• Unity Catalog関数を変わらずツールとして使えるのか

検証はDatabricks on AWS上でノートブックを作成して行いました。
クラスタはサーバレスを利用しています。

Step0. 準備

ノートブックを作成し、必要なパッケージをインストールします。
langgraphに加えてdatabricks-langchainなどDatabricksのLLMやツールを利用するパッケージをインストールし、またMLflowの最新化を行っています。

%pip install langgraph databricks-langchain unitycatalog-ai[databricks] unitycatalog-langchain[databricks]
%pip install mlflow-skinny[databricks]

%restart_python

Step1. ツールの作成

DatabricksのUnity Catalog関数がFunctional API下でも問題なく使えるか確認したかったので、確認用にTavilyを使ったWeb検索ツールを作成します。

外部APIを呼び出すツールの作成については下記の公式ドキュメントを参照ください。

まず、Tavilyへの接続オブジェクトを作成。
bearer_tokenは事前にDatabricks Secretへ保管しておいたTavilyのAPIキーを設定しています。

%sql
CREATE CONNECTION tavily_api TYPE HTTP
OPTIONS (
  host 'https://api.tavily.com',
  port '443',
  base_path '/',
  bearer_token secret ('tavily','APIトークンのキー')
);

作成した接続オブジェクトを用いて、Unity Catalog関数を作成します。

%sql
CREATE OR REPLACE FUNCTION training.llm.tavily_web_search(
  query STRING COMMENT 'query text'
)
RETURNS STRING
COMMENT 'Performs a web search using Tavily'
RETURN (SELECT http_request(
  conn => 'tavily_api',
  method => 'POST',
  path => 'search',
  json => to_json(named_struct(
    'query', query
  )),
  headers => map(
    'Content-Type', "application/json"
  )
)).text

Step2. エージェントをFunctional APIを使って作成

では、チュートリアルを基にエージェントを作成しましょう。

まずは利用するLLMやツールの取得。

from langchain_core.tools import tool
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit
from unitycatalog.ai.core.databricks import DatabricksFunctionClient
from databricks_langchain import ChatDatabricks

client = DatabricksFunctionClient()

# Databricks上で提供されているPay-per-tokenのモデルを利用
model = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")

# UC関数を使用してツールキットを作成
func_name = "training.llm.tavily_web_search"
uc_toolkit = UCFunctionToolkit(function_names=[func_name], client=client)
tools = uc_toolkit.tools

次に実行するTaskを定義します。
Functional APIでは、いわゆるノードをタスクとして定義します。
関数に@taskデコレータを付けることでタスクの宣言となります。
また、MLflow Tracingが使えるか確認するために、各タスク関数には@mlflow.traceデコレータも付けました。

各タスクはそれぞれLLMへのクエリ実行(Tool Calling含む)と、ツールの実行を担っています。

from langchain_core.messages import ToolMessage
from langgraph.func import entrypoint, task
import mlflow

tools_by_name = {tool.name: tool for tool in tools}

@task
@mlflow.trace
def call_model(messages):
    """Call model with a sequence of messages."""
    response = model.bind_tools(tools).invoke(messages)
    return response

@task
@mlflow.trace
def call_tool(tool_call):
    tool = tools_by_name[tool_call["name"]]
    observation = tool.invoke(tool_call["args"])
    return ToolMessage(content=observation, tool_call_id=tool_call["id"])

最後にエージェントのワークフローを定義します。
従来のGraph APIだとノードやエッジの定義を行っていましたが、Functional APIでは通常の関数を使用してワークフローを定義します。
したがって、条件分岐やループも通常の記法で制御できます。

ワークフローの定義関数には@entrypointを付与する必要があります。
また、こちらにもMLflow Tracingが効くのか確認するために@mlflow.traceデコレータを付与しています。

from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver

checkpointer = MemorySaver()

@entrypoint(checkpointer=checkpointer)
@mlflow.trace(span_type="AGENT")
def agent(messages):
    # まず、モデルに問い合わせ
    llm_response = call_model(messages).result()
    while True:

        # 問い合わせ結果がツール呼び出しではない場合、終了
        if not llm_response.tool_calls:
            break

        # ツールの実行
        tool_result_futures = [
            call_tool(tool_call) for tool_call in llm_response.tool_calls
        ]
        tool_results = [fut.result() for fut in tool_result_futures]

        # ツールの実行結果をmessagesに追加
        messages = add_messages(messages, [llm_response, *tool_results])

        # messagesを使って再度モデルに問い合わせ
        llm_response = call_model(messages).result()

    return llm_response

ループ制御が通常のPython処理として書けるので、慣れた形で定義できますね。
（処理の見通しがいいかどうかは賛否ありそうですが。。。）

Step3. 実行

では、定義したエージェントワークフローを実行してみます。
Web検索ツールを使うので、時事系の質問を入れてみます。

config = {"configurable": {"thread_id": "1"}}

user_message = {"role": "user", "content": "現在の日本の総理大臣は？"}
print(user_message)

for step in agent.stream([user_message], config):
    for task_name, message in step.items():
        if task_name == "agent":
            continue  # Just print task updates
        print(f"\n{task_name}:")
        message.pretty_print()

出力結果

{'role': 'user', 'content': '現在の日本の総理大臣は？'}

call_model:
================================== Ai Message ==================================
Tool Calls:
  training__llm__tavily_web_search (call_f6c08cba-f63d-4158-b9a8-8b7f9ff08d55)
 Call ID: call_f6c08cba-f63d-4158-b9a8-8b7f9ff08d55
  Args:
    query: current Japanese prime minister

call_tool:
================================= Tool Message =================================

{"format": "SCALAR", "value": "{\"query\":\"current Japanese prime minister\",\"follow_up_questions\":null,\"answer\":null,\"images\":[],\"results\":[{\"title\":\"List of Japan Prime Ministers (1885-2024) - Jagran Josh\",\"url\":\"https://www.jagranjosh.com/general-knowledge/list-of-japan-prime-ministers-1632984150-1\",\"content\":\"Japan's current Prime Minister is Shigeru Ishiba, who took office on October 1, 2024. He succeeded Fumio Kishida, who served from October 4, 2021, until November 10, 2021. Early Life.\",\"score\":0.9323611,\"raw_content\":null},{\"title\":\"Japan's parliament elects Shigeru Ishiba as prime minister\",\"url\":\"https://www.npr.org/2024/10/01/nx-s1-5134532/shigeru-ishiba-becomes-japan-prime-minister\",\"content\":\"Japan\u2019s parliament elects Shigeru Ishiba as prime minister TOKYO \u2014 Japan\u2019s parliament on Tuesday formally elected Shigeru Ishiba, head of the governing Liberal Democratic Party, as the country\u2019s new prime minister. Japan's new PM promises to bring continuity and changes to dealings with U.S. Opposition parties criticized Ishiba for allowing only a short period of time for his policies to be examined and discussed in parliament before the national election. The majority of his Cabinet ministers, like Ishiba, are expected to be unaffiliated with factions led and controlled by party heavyweights, and none are from former Prime Minister Shinzo Abe's powerful group linked to damaging scandals. Ishiba, first elected to parliament in 1986, has served as defense minister, agriculture minister and in other key Cabinet posts, and was LDP secretary general under Abe. Year of Global Elections\",\"score\":0.85509074,\"raw_content\":null},{\"title\":\"Shigeru Ishiba elected Japan's prime minister by Parliament after ...\",\"url\":\"https://www.cbsnews.com/news/shigeru-ishiba-japan-prime-minister-parliament-fumio-kishida-administration-scandals/\",\"content\":\"Shigeru Ishiba elected Japan's prime minister by Parliament after predecessor's administration was rocked by scandals Tokyo \u2014 Japan's parliament formally elected Shigeru Ishiba, head of the governing Liberal Democratic Party, as the country's new prime minister Tuesday. The majority of his Cabinet ministers are expected, as is Ishiba, to be unaffiliated with factions led and controlled by party heavyweights, and none are from former Prime Minister Shinzo Abe's powerful group, which has been linked to damaging scandals. Ishiba, first elected to parliament in 1986, has served as defense minister, agriculture minister and in other key Cabinet posts, and was LDP secretary general under Abe. More from CBS News\",\"score\":0.84707206,\"raw_content\":null},{\"title\":\"Japan parliament elects Shigeru Ishiba as the new prime minister to ...\",\"url\":\"https://apnews.com/article/japan-cabinet-kishida-ishiba-fa7f80b614a8bf298852b6766dbf6403\",\"content\":\"Japan parliament elects Shigeru Ishiba as the new prime minister to replace Kishida | AP News Movie Reviews What to Stream Television Book Reviews Music Celebrity Interviews Back to school Food & Recipes Gardening Homes Travel Fashion Pets 2 of 13\u00a0|\u00a0Japan\u2019s newly elected Prime Minister Shigeru Ishiba speaks during a press conference at the prime minister\u2019s office in Tokyo Tuesday, Oct. 1, 2024. 3 of 13\u00a0|\u00a0Japan\u2019s newly elected Prime Minister Shigeru Ishiba speaks during a press conference at the prime minister\u2019s office in Tokyo Tuesday, Oct. 1, 2024. 4 of 13\u00a0|\u00a0Japan\u2019s newly elected Prime Minister Shigeru Ishiba speaks during a press conference at the prime minister\u2019s office in Tokyo Tuesday, Oct. 1, 2024.\",\"score\":0.8245895,\"raw_content\":null},{\"title\":\"Japan parliament reelects Prime Minister Ishiba | AP News\",\"url\":\"https://apnews.com/article/japan-politics-ishiba-cabinet-5945feb06a07730a98ebb3e0b46905bd\",\"content\":\"Movie reviews What to Stream Television Book reviews Music Celebrity Interviews 11 of 19\u00a0|\u00a0A staff member at parliament puts up two names, Shigeru Ishiba (of Japanese Prime Minister), right, and Yoshihiko Noda (leader of the Constitutional Democratic Party) for a runoff vote for a new prime minister at a special parliamentary session of the lower house Monday, Nov. 11, 2024, in Tokyo. 11 of 19\u00a0|\u00a0A staff member at parliament puts up two names, Shigeru Ishiba (of Japanese Prime Minister), right, and Yoshihiko Noda (leader of the Constitutional Democratic Party) for a runoff vote for a new prime minister at a special parliamentary session of the lower house Monday, Nov. 11, 2024, in Tokyo.\",\"score\":0.81875163,\"raw_content\":null}],\"response_time\":0.07}"}

call_model:
================================== Ai Message ==================================

現在の日本の総理大臣は、石破茂（しげる）です。

ちゃんとUnity Catalogに登録されているWeb検索ツールを実行し、正しく結果が得られました。
また、MLflow Tracingも有効になり、エントリーポイントとタスクの結果が期待通り記録されました。

LLMに問い合わせ(ツール選択)→ツール実行→LLMに(ツール実行結果を含めて)再度問い合わせ　という流れになっているのが確認できます。

まとめ

LangGraphのFunctional APIをDatabricksの特徴的な機能と組み合わせて試してみました。

Functional APIに関する個人的な感想としては、条件分岐やループについてはこれまでのGraph APIより書きやすい印象です。まだ少ししか書いてないのでより複雑な処理を記載していくとどうなるかわかりませんが。

一方でグラフを可視化する機能が(まだ)実装されていないので、グラフ全体の流れを把握したりすることが難しくなりました。このあたりもあってGraph APIと補完しながら使うことがいいのだろうと思います。
正直、まだどういったケースでGraph APIが良いのか・Functional APIが良いのか、という見極めができていないのでこのあたりは習熟が必要そうです。

また、最初Functional APIはあくまでおまけ程度の実装なのかなーと思っていたのですがドキュメントを読んでいる分だと結構大きい規模のAPIであり、Graph APIとFunctional APIの2本柱で今後展開を進めていく雰囲気を感じています。
LangGraphは個人的に好きなパッケージなのでこの先もどんどん発展していくことを期待しています。

Unity Catalog関数と組み合わせてエージェント実装するの、ますます楽しくなるなあ。