LangGraphで AIエージェントをまなんでいく - その4 一連のプロセスを繰り返す-

from langchain_google_genai import ChatGoogleGenerativeAI

def multiply(a: int, b: int) -> int:
    """Multiply a and b.

    Args:
        a: first int
        b: second int
    """
    return a * b

# This will be a tool
def add(a: int, b: int) -> int:
    """Adds a and b.

    Args:
        a: first int
        b: second int
    """
    return a + b

def divide(a: int, b: int) -> float:
    """Divide a and b.

    Args:
        a: first int
        b: second int
    """
    return a / b

tools = [add, multiply, divide]
llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
llm_with_tools = llm.bind_tools(tools)

from langgraph.graph import MessagesState
from langchain_core.messages import HumanMessage, SystemMessage

# System message
sys_msg = SystemMessage(content="あなたは、一連の入力に対して算術演算を実行するタスクを与えられた親切なアシスタントでぇす。")

# Node
def assistant(state: MessagesState):
   return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}

from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from IPython.display import Image, display

# Graph
builder = StateGraph(MessagesState)

# ノードの定義。このノードが実行されます
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))

# エッジの定義：制御フローがどのように動くかを決める。
builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    # アシスタントからの最新のメッセージ（結果）がツールコールである場合 -> tools_conditionはtoolsにルーティングする。
    # アシスタントからの最新のメッセージ（結果）がツールコールでない場合 -> tools_conditionはENDにルーティングする。
    tools_condition,
)
builder.add_edge("tools", "assistant")
react_graph = builder.compile()

# Show
display(Image(react_graph.get_graph(xray=True).draw_mermaid_png()))

messages = [HumanMessage(content="3と4を足す。その出力に2を掛ける。 さらにその出力を5で割る。")]
messages = react_graph.invoke({"messages": messages})
for m in messages['messages']:
    m.pretty_print()

messages = [HumanMessage(content="大当たり確率が1/199の台で、100回転以内に当たりが引ける期待値は？")]
messages = react_graph.invoke({"messages": messages})