DatabricksでLangGraphのクイックスタートを動かしてみる(その5)

こちらの続きです。

https://qiita.com/taka_yayoi/items/e2a7fcafd21acd31eeff

こちらのPart 5: Customizing Stateを動かします。

https://langchain-ai.github.io/langgraph/tutorials/introduction/#part-5-customizing-state

# パート5: 状態のカスタマイズ

これまで、メッセージのリストという1つのエントリを持つシンプルな状態に依存してきました。このシンプルな状態でも多くのことができますが、メッセージリストに依存せずに複雑な動作を定義したい場合は、状態に追加のフィールドを追加できます。

## セットアップ

まず、必要なパッケージをインストールし、環境を設定します：

```py
%%capture --no-stderr
%pip install -U langgraph langsmith langchain_openai openai tavily-python langchain_community
%restart_python
```

```py
import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(scope="demo-token-takaaki.yayoi", key="openai_api_key")

# TavilyのAPIキー
os.environ["TAVILY_API_KEY"] = "TavilyのAPIキー"
```

ここでは、チャットボットが特定の情報を検索ツールを使って見つけ、それを人間に転送する新しいシナリオを示します。エンティティの誕生日を調査するチャットボットを作成しましょう。状態に`name`と`birthday`のキーを追加します：

```py
from typing import Annotated

from typing_extensions import TypedDict

from langgraph.graph.message import add_messages


class State(TypedDict):
    messages: Annotated[list, add_messages]
    name: str
    birthday: str
```

この情報を状態に追加することで、他のグラフノード（例：情報を保存または処理する下流ノード）やグラフの永続化レイヤーから簡単にアクセスできるようになります。

ここでは、`human_assistance`ツール内で状態キーを設定します。これにより、人間が情報を状態に保存する前にレビューすることができます。今回は`Command`を使用して、ツール内から状態更新を発行します。`Command`の使用例については[こちら](https://langchain-ai.github.io/langgraph/concepts/low_level/#using-inside-tools)を参照してください。

```py
from langchain_core.messages import ToolMessage
from langchain_core.tools import InjectedToolCallId, tool

from langgraph.types import Command, interrupt


@tool
# 状態更新のためにToolMessageを生成しているため、一般的に対応するツール呼び出しのIDが必要です。
# LangChainのInjectedToolCallIdを使用して、この引数がツールのスキーマでモデルに公開されないことを示すことができます。
def human_assistance(
    name: str, birthday: str, tool_call_id: Annotated[str, InjectedToolCallId]
) -> str:
    """人間のアシストをリクエスト"""
    human_response = interrupt(
        {
            "question": "これは正しいですか？",
            "name": name,
            "birthday": birthday,
        },
    )
    # 情報が正しい場合、そのまま状態を更新します。
    if human_response.get("correct", "").lower().startswith("y"):
        verified_name = name
        verified_birthday = birthday
        response = "正しい"
    # そうでない場合、人間のレビュアーから情報を受け取ります。
    else:
        verified_name = human_response.get("name", name)
        verified_birthday = human_response.get("birthday", birthday)
        response = f"修正を行いました: {human_response}"

    # 今回はツール内でToolMessageを使用して明示的に状態を更新します。
    state_update = {
        "name": verified_name,
        "birthday": verified_birthday,
        "messages": [ToolMessage(response, tool_call_id=tool_call_id)],
    }
    # ツール内で状態を更新するためにCommandオブジェクトを返します。
    return Command(update=state_update)
```

それ以外は、グラフの残りの部分は同じです：

```py
from typing import Annotated

from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode, tools_condition

tool = TavilySearchResults(max_results=2)
tools = [tool, human_assistance]
llm = ChatOpenAI(model="gpt-4o-mini")
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State):
    message = llm_with_tools.invoke(state["messages"])
    assert len(message.tool_calls) <= 1
    return {"messages": [message]}

graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot)

tool_node = ToolNode(tools=tools)
graph_builder.add_node("tools", tool_node)

# 条件付きエッジを追加
graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# ツールが呼び出されるたびに、次のステップを決定するためにチャットボットに戻る
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")

memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)
```

アプリケーションにLangGraphライブラリの「誕生日」を調べるように促しましょう。必要な情報が得られたら、チャットボットに`human_assistance`ツールに連絡するよう指示します。ツールの引数に`name`と`birthday`を設定することで、チャットボットがこれらのフィールドの提案を生成するように強制します。

```py
user_input = (
    "LangGraphがリリースされた時期を調べてください。"
    "答えがわかったら、human_assistanceツールを使ってレビューしてください。"
)
config = {"configurable": {"thread_id": "1"}}

events = graph.stream(
    {"messages": [{"role": "user", "content": user_input}]},
    config,
    stream_mode="values",
)
for event in events:
    if "messages" in event:
        event["messages"][-1].pretty_print()
```

![Screenshot 2025-02-04 at 7.46.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/85f2668f-6863-887e-f09c-6f12f2ecbf47.png)

再び`human_assistance`ツールで`interrupt`が発生しました。この場合、チャットボットが正しい日付を特定できなかったため、私たちが提供します：

```py
human_command = Command(
    resume={
        "name": "LangGraph",
        "birthday": "Jan 17, 2024",
    },
)

events = graph.stream(human_command, config, stream_mode="values")
for event in events:
    if "messages" in event:
        event["messages"][-1].pretty_print()
```

![Screenshot 2025-02-04 at 7.47.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/42ea65d5-3ad4-2803-71e0-dc5db6352c0b.png)
![Screenshot 2025-02-04 at 7.47.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/89a0325a-c17e-099b-4dc1-a1205af7187b.png)

これらのフィールドが状態に反映されていることに注意してください：

```py
snapshot = graph.get_state(config)

{k: v for k, v in snapshot.values.items() if k in ("name", "birthday")}
```
```
{'name': 'LangGraph', 'birthday': 'Jan 17, 2024'}
```

これにより、下流ノード(例：情報をさらに処理または保存するノード)から簡単にアクセスできるようになります。

## 状態の手動更新

LangGraphはアプリケーションの状態に対して高い制御を提供します。例えば、任意の時点(中断時を含む)で、`graph.update_state`を使用してキーを手動で上書きすることができます。

```py
graph.update_state(config, {"name": "LangGraph (library)"})
```
![Screenshot 2025-02-04 at 7.48.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/4ef1aa62-90c0-dc05-ba41-944d4300a5e1.png)

`graph.get_state`を呼び出すと、新しい値が反映されていることがわかります：

```py
snapshot = graph.get_state(config)

{k: v for k, v in snapshot.values.items() if k in ("name", "birthday")}
```
```
{'name': 'LangGraph (library)', 'birthday': 'Jan 17, 2024'}
```

手動の状態更新でもLangSmithに[トレースが生成](https://smith.langchain.com/public/7ebb7827-378d-49fe-9f6c-5df0e90086c8/r)されます。必要に応じて、[このガイド](https://langchain-ai.github.io/langgraph/how-tos/human_in_the_loop/edit-graph-state/)で説明されているように、人間が介在するワークフローを制御するためにも使用できます。代わりに`interrupt`機能の使用が一般的に推奨されます。これは、状態更新とは独立して、人間が介在するインタラクションでデータを送信できるためです。

**おめでとうございます！** より複雑なワークフローを促進するために状態にカスタムキーを追加し、ツール内から状態更新を生成する方法を学びました。

チュートリアルはほぼ完了ですが、`checkpointing`と`state updates`を結びつける前に確認したいもう一つの概念があります。

次で最後です。

https://qiita.com/taka_yayoi/items/d0d264b50128aedd7e90

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)