LangGraphの使い方(簡易AIエージェント(Gemini)編)

# はじめに
[LangGraphの使い方(分岐編)](https://qiita.com/te_yama/items/02df3b4216d47b0ebbbf)の続きになります。
前回は分岐を行うグラフを定義しました。
今回はGoogleのGeminiを利用したAIエージェントを定義したいと思います。

# AIエージェント
## インストール
以下のライブラリをインストールします。
```bash
pip install langgraph langchain-google-genai
```

## LLM(Gemini)の定義
Geminiの場合は`ChatGoogleGenerativeAI`を使用して、LLMを定義します。

https://python.langchain.com/api_reference/google_genai/chat_models/langchain_google_genai.chat_models.ChatGoogleGenerativeAI.html#

`GOOGLE_API_KEY`を環境変数に設定しておく必要があります。
```python
import os
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY
```

そして、Geminiのモデル名を指定してLLMを定義します。
今回はGemini 2.0 Flash(試用版)を使用します。
```python
from langchain_google_genai import ChatGoogleGenerativeAI

model = ChatGoogleGenerativeAI(model='gemini-2.0-flash-exp', temperature=0)
```

## Tool Callingの定義
Tool Callingとは、LLMから外部機能(関数)を呼び出す機能のことです。
(よく知られているFunction callingと同等のものだと思っています。)

LangGraphでは、`Pydantic`を使用してTool Callingを定義します。
今回は2つのtoolを定義します。
その際、docstringで各Toolに対してどのようなToolなのかを記載する。

- GetWeather
    - 指定された場所の天気を取得するTool
- GetPopulation
    - 指定された場所の人口を取得するTool

```python
from pydantic import BaseModel, Field

class GetWeather(BaseModel):
    '''指定された場所の現在の天気を取得する'''
    location: str = Field(
        ..., description="都市の名前"
    )

class GetPopulation(BaseModel):
    '''指定された場所の現在の人口を取得する'''
    location: str = Field(
        ..., description="都市の名前"
    )
```

定義したToolをLLMに`bind_tools`でバインドする。
```python
tools = [GetWeather, GetPopulation]
model_with_tools = model.bind_tools(tools)
```

### テスト
ToolをバインドしたLLMに対して、`東京の人口は？`と質問する。
そのレスポンスの`tool_calls`に、質問の内容に適した定義したToolが選択される。

```python
response = model_with_tools.invoke("東京の人口は？")

response.tool_calls
[{'name': 'GetPopulation',
  'args': {'location': '東京'},
  'id': '0b28683f-1621-42ea-9f5e-d9d89a5e1ba4',
  'type': 'tool_call'}]
```
人口について質問したので、`GetPopulation`が選択されて、引数として都市の名前である`東京`が選択されていることが確認できる。

## グラフの定義
### Stateの定義
`messages`を状態として持っておくStateを定義する。

```python
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

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

### Agent Nodeの定義
ToolをバインドしたLLMをエージェントとして、Nodeを定義する。
ユーザーからの質問を入力して、返ってきた内容を`messages`に追加する。

```python
def agent(state: State):
    state["messages"].append(model_with_tools.invoke(state["messages"]))
    return state
```

### Tool Nodeの定義
定義したToolに対応する関数を定義し、`tool_by_name`のDictとして紐づけを行う。
Stateの`messages`内の最後のデータを見て、`tool_calls`の`name`からToolの名前を取得し、実行すべき関数を決める。
その後、引数を取得しその関数を実行し得られたデータをStateの`messages`に格納する。

本来であれば、APIを呼び出して最新の情報などを取得するのが一般的ですが、今回は簡易的に返すデータは引数として選ばれた都市の名前が変わる固定文字列としている。

```python
from langchain_core.tools import tool
from langchain_core.messages import  ToolMessage

def tools(state: State):
    tool_by_name = {
        "GetWeather": get_weather,
        "GetPopulation": get_population
        }
    last_message = state["messages"][-1]
    
    tool_function = tool_by_name[last_message.tool_calls[0]["name"]]
    tool_output = tool_function.invoke(
        last_message.tool_calls[0]["args"]
        )

    state["messages"].append(ToolMessage(
        content=tool_output, 
        tool_call_id=last_message.tool_calls[0]["id"]
        )
    )
    return state

@tool
def get_weather(location: str):
    """天気に関する情報を返す関数"""
    return f"{location}は晴れです。"

@tool
def get_population(location: str):
    """人口に関する情報を返す関数"""
    return f"{location}の人口は100人です。"
```

## 条件付きEdgeの定義
Stateの`messages`内の最後のデータに`tool_calls`があるかないかでToolを実行するか、終わるかを条件としている。

```python
from typing import Literal
from langgraph.graph import START ,END

def should_continue(state: State) -> Literal["tools", END]:
    messages = state['messages']
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return END
```

## StateGraphの定義
上記で定義したものを使用して、StateGraphを定義する。
```python
from langgraph.graph import StateGraph
from IPython.display import Image, display

graph = StateGraph(State)

graph.add_node("agent", agent)
graph.add_node("tools", tools)

graph.add_edge("tools", 'agent')

graph.set_entry_point("agent")

graph.add_conditional_edges(
    "agent",
    should_continue
)

graph = graph.compile()
display(Image(graph.get_graph().draw_mermaid_png()))
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3968617/cae59175-9126-8d58-4aab-1e73fc7236f8.png)

<details><summary>全体のコード</summary>

```python
import os
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

from langchain_google_genai import ChatGoogleGenerativeAI

model = ChatGoogleGenerativeAI(model='gemini-2.0-flash-exp', temperature=0)

from pydantic import BaseModel, Field
class GetWeather(BaseModel):
    '''指定された場所の現在の天気を取得する'''

    location: str = Field(
        ..., description="都市の名前"
    )


class GetPopulation(BaseModel):
    '''指定された場所の現在の人口を取得する'''

    location: str = Field(
        ..., description="都市の名前"
    )

tools = [GetWeather, GetPopulation]
model_with_tools = model.bind_tools(tools)

from typing import Annotated, Literal
from typing_extensions import TypedDict
from langgraph.graph import START ,END
from langgraph.graph import StateGraph
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from IPython.display import Image, display

# Stateを宣言
class State(TypedDict):
    messages: Annotated[list, add_messages]

# Nodeを宣言
def agent(state: State):
    state["messages"].append(model_with_tools.invoke(state["messages"]))
    return state

def should_continue(state: State) -> Literal["tools", END]:
    messages = state['messages']
    last_message = messages[-1]
    # If the LLM makes a tool call, then we route to the "tools" node
    if last_message.tool_calls:
        return "tools"
    # Otherwise, we stop (reply to the user)
    return END

def tools(state: State):
    tool_by_name = {
        "GetWeather": get_weather,
        "GetPopulation": get_population
        }
    last_message = state["messages"][-1]
    tool_function = tool_by_name[last_message.tool_calls[0]["name"]]
    tool_output = tool_function.invoke(
        last_message.tool_calls[0]["args"]
        )
    state["messages"].append(ToolMessage(
        content=tool_output, 
        tool_call_id=last_message.tool_calls[0]["id"]
        )
    )
    return state

@tool
def get_weather(location: str):
    """天気に関する情報を返す関数"""
    return f"{location}は晴れです。"

@tool
def get_population(location: str):
    """人口に関する情報を返す関数"""
    return f"{location}の人口は100人です。"

# Graphの作成
graph = StateGraph(State)

# Nodeの追加
graph.add_node("agent", agent)
graph.add_node("tools", tools)

graph.add_edge("tools", 'agent')

# Graphの始点を宣言
graph.set_entry_point("agent")

graph.add_conditional_edges(
    "agent",
    should_continue
)

# Graphをコンパイル
graph = graph.compile()
display(Image(graph.get_graph().draw_mermaid_png()))

```

</details>


## 実行
### 実行1
`東京の天気は？`と質問をする。
```python
graph.invoke({"messages": "東京の天気は？"}, debug=True)
```

出力(省略)
```bash
 AIMessage(content='東京は晴れです。\n', ～
```

### 実行2
`東京の人口は？`と質問をする。
```python
graph.invoke({"messages": "東京の人口は？"}, debug=True)
```

出力(省略)
```bash
 AIMessage(content='東京の人口は100人です。\n', ～
```

質問の内容の意図に応じたToolが選択されて、定義した関数が実行されていることが確認できました！

### 実行3
少しニュアンスの異なる聞き方をしてみたいと思います。
`ベルリンの今の人の数は？`と質問をする。
```python
graph.invoke({"messages": "ベルリンの今の人の数は？"}, debug=True)
```

出力(省略)
```bash
 AIMessage(content='ベルリンの人口は100人です。\n', ～
```

人口というワードが無くても、人の数という似た意味だと捉えてToolを選択してくれました！


# おわりに
今回はGeminiを利用した簡易AIエージェントを作成してみました。
質問の内容に合った定義したToolを選択されて、その関数が実行されることを確認しました。
入力データに基づいた意思決定を行うAIエージェントは、今後の生成AIの流行になると言われているので、基礎技術はしっかりと理解していきたいと思いました。
今後はもう少し複雑なAIエージェントを作ってみたい。