[LangGraph] ToolNode

LLM에 Tool을 binding 해서 LLM이 tool_calls를 생성했을 때, 적절한 arguments를 사용해 해당 tool을 실행하도록 하는 ToolNode에 대해 자세히 알아보자. 방금 작성한 그대로 로직을 작성할 수도 있지만(참고 - [LangGrpah] Tool Binding), LangGraph는 ToolNode를 사전 정의(pre-built)해서 제공한다.

 

내부적으로는 LLM에게 tool의 목록을 전달하고 (bind_tools), LLM이 사용자의 질문을 기반으로 tool 실행이 필요하다고 판단하면 해당 tool의 이름과 arguments를 반환한다. 그러면 해당 tool과 arguments로 함수를 실행하게 된다. 이때 tool list를 갖고, LLM이 반환한 tool_calls를 기반으로 함수를 실행할 수 있도록 노드로 구현한 것이 ToolNode다.

 

도구 정의

먼저 파이썬 코드를 실행하는 execute_python과 location의 따라 서로 다른 문자열을 반환하는get_weather를 정의하자. 그리고 이 두 함수를 tool 콜백 함수를 사용해서 llm에 binding 하기 좋은 형태로 변환하자.

다음으로, 두 개의 tool을 list 형태로 만든 뒤, ToolNode를 초기화하자.

from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langchain_experimental.tools.python.tool import PythonAstREPLTool
from langgraph.prebuilt import ToolNode

@tool
def execute_python(code: str):
    """Call to excute python code."""
    return PythonAstREPLTool().invoke(code)

@tool
def get_weather(location: str):
    """Call to get the current weather."""
    if location.lower() in ["seoul", "busan"]:
        return "The temperature is 5 degrees and it's cloudy."
    else:
        return "The temperature is 30 degrees and it's sunny."

# Tool Node 초기화
tools = [execute_python, get_weather]
tool_node = ToolNode(tools=tools)

ToolNode 수동 호출해보기

ToolNode는 State의 messages list의 마지막 메세지가 tool_calls 인자가 있는지 없는지 여부로 tool을 호출할지를 결정한다.

 

일반적으로는 AIMessage를 수동으로 생성하지 않고, LangChain의 LLM이 생성하지만, 먼저 ToolNode를 수동으로 호출해 보자.

AIMessage에서 content는 보통 빈 문자열이 들어가고, tool_calls 속성에 List[Dict] 형태로, 호출할 도구의 이름, 인자, ID, 유형의 키 키값으로 들어간다.

message_with_single_tool_call = AIMessage(
    content="",
    tool_calls=[
        {
            "name": "get_weather",
            "args": {"location": "seoul"},
            "id": "tool_call_id",
            "type": "tool_call"
        }
    ]
)

print(tool_node.invoke({"messages": [message_with_single_tool_call]}))

 

ToolNode는 execute_python과 get_weather 두 개의 tool을 갖고 있지만, invoke 한 결과는 ToolMessage로 반환되며, content에는 get_weather(location="seoul") 을 실행한 결과가 들어간다.

{'messages': [ToolMessage(content="The temperature is 5 degrees and it's cloudy.", name='get_weather', tool_call_id='tool_call_id')]}

병렬 수행

AIMessage의 tool_calls 인자에 리스트로 여러 tool을 전달하면, ToolNode가 병렬적으로 도구 호출을 수행한다. 아래는 execute_python과 get_weather 두 가지 tool_call을 tool_calls 리스트에 전달했다. 그리고 결과는 예상대로 2개의 ToolMessage가 반환된다.

message_with_multiple_tool_call = AIMessage(
    content="",
    tool_calls=[
        {
            "name": "get_weather",
            "args": {"location": "busan"},
            "id": "tool_call_id_1",
            "type": "tool_call"
        },
        {
            "name": "execute_python",
            "args": {"code": "3 + 3"},
            "id": "tool_call_id_2",
            "type": "tool_call"
        }
    ]
)

print(tool_node.invoke({"messages": [message_with_multiple_tool_call]}))

{'messages': [ToolMessage(content="The temperature is 5 degrees and it's cloudy.", name='get_weather', tool_call_id='tool_call_id_1'), ToolMessage(content='6', name='execute_python', tool_call_id='tool_call_id_2')]}

ChatModel과 함께 사용

앞서 말한 것처럼, 일반적으로는 AIMessage를 수동으로 생성하지 않고, LangChain의 LLM이 생성한다. bind_tools 메소드를 호출해서 LLM에 도구를 인식시켜 줄 수 있다.

실제로는 아래와 같은 프롬프트가 들어간다! (참고)

"""You are an assistant that has access to the following set of tools. Here are the names and descriptions for each tool:

{rendered_tools}

Given the user input, return the name and input of the tool to use. Return your response as a JSON blob with 'name' and 'arguments' keys."""

 

그럼 이번에는 tool을 들고 있는 llm_with_tools를 정의하고, 서울의 현재 날씨를 물어보자. 그러면 LLM이 get_weather tool을 호출해야 한다고 판단하고, AIMessage에 tool_calls 인자가 채워져 반환된다.

from langchain_openai import ChatOpenAI

llm_with_tools = ChatOpenAI(model="gpt-4o-mini").bind_tools(tools=tools)
llm_with_tools.invoke("seoul의 현재 날씨는 어때요?").tool_calls

[{'name': 'get_weather', 
  'args': {'location': 'seoul'}, 
  'id': 'call_7TfQhajj4O3fM1FROdH4gAwC', 
  'type': 'tool_call'}]

따라서, 해당 반환값을 직접 ToolNode에 전달해 실행할 수 있다.

tool_node.invoke({
    "messages": [llm_with_tools.invoke("seoul의 현재 날씨는 어때요?")]
})

{'messages': [ToolMessage(content="The temperature is 5 degrees and it's cloudy.", name='get_weather', tool_call_id='call_jDy3hBVkBKHzHDkzC1UuzyOX')]}

ReACT 에이전트

다음으로 그래프에서 ToolNode를 통합하여 사용하는 방법을 살펴보자. 여기서는 ReAct 에이전트를 구현해 보자. 이 에이전트는 일부 쿼리를 input으로 사용하고, 쿼리에 적절한 답변을 생성하기에 충분한 정보가 수집될 때까지 tool을 반복적으로 호출한다.

이때, RemoveMessage를 사용해서 최근 5개의 대화까지만 사용하도록 하자.

from typing import Annotated, TypedDict, Literal
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages, RemoveMessage
from langgraph.checkpoint.memory import MemorySaver


# 상태 정의
class State(TypedDict):
    messages: Annotated[list, add_messages]

# conditional edge에 사용할 함수 정의
def should_continue(state: State) -> Literal["tools", "delete_message"]:
    last_message = state["messages"][-1]
    if last_message.tool_calls:    # 마지막 메세지에 tool_calls가 있다면
        return "tools"
    return "delete_message"

# chat_model Node 정의
def chat_model(state: State):
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# delete_message Node 정의
def delete_message(state: State):
    """메세지 개수가 5개가 넘어갈 경우, 오래된 메세지 삭제해서 최신 5개만 유지"""
    if len(state["messages"]) > 5:
        return {"messages": [RemoveMessage(id=message.id) for message in state["messages"][:-5]]}

flow = StateGraph(State)
flow.add_node("chat_model", chat_model)
flow.add_node("delete_message", delete_message)
flow.add_node("tools", tool_node)

flow.add_edge(START, "chat_model")
flow.add_edge("tools", "chat_model")
flow.add_edge("delete_message", END)

flow.add_conditional_edges(
    "chat_model",
    should_continue
)

memory = MemorySaver()
graph = flow.compile(checkpointer=memory)

 

만들어진 그래프를 시각화해 보자.

from IPython.display import display, Image

display(Image(graph.get_graph().draw_mermaid_png()))

 

 

 

사용자 질문이 들어오면 execute_python와 get_weather 2개의 도구를 들고 있는 chat_model 노드에서 2가지 액션 중 하나를 취할 수 있다.

1. tool 호출
   1. 파이썬 함수 실행이 필요한 경우 execute_python
   2. 현재 날씨가 필요한 경우 get_weather
2. 직접 답변
   1. 굳이 Tool 호출이 필요하지 않은 경우 직접 답변

그리고 최종적으로 Tool 호출이 일어나지 않고 답변이 완료됐다면, delete_message 노드로 이동해서 가장 최근 5개의 메세지 이력만 남기고 나머지는 지운다.

 

테스트해 보자. 먼저, 파이썬 코드를 작성하도록 요청해 보자. (주피터 노트북 환경에서 실행하면 그래프 결과도 확인 가능하다)

input = {"messages": [("user", "y=2x+3 그래프를 파이썬으로 그려주세요.")]}
config = {"configurable": {"thread_id": "a"}}
events = graph.stream(input, config, stream_mode="values")

for event in events:
    event["messages"][-1].pretty_print()

================================ Human Message =================================

y=2x+3 그래프를 파이썬으로 그려주세요.

================================== Ai Message ==================================
Tool Calls:
  execute_python (call_dyChs2ZSWFvXQbWv1W4DQUvZ)
 Call ID: call_dyChs2ZSWFvXQbWv1W4DQUvZ
  Args:
    code: import matplotlib.pyplot as plt
import numpy as np

# Define the function
def linear_function(x):
    return 2 * x + 3

# Generate x values
x_values = np.linspace(-10, 10, 400)
# Generate y values based on the linear function
y_values = linear_function(x_values)

# Create the plot
plt.figure(figsize=(10, 6))
plt.plot(x_values, y_values, label='y = 2x + 3', color='blue')
plt.title('Graph of y = 2x + 3')
plt.xlabel('x')
plt.ylabel('y')
plt.axhline(0, color='black',linewidth=0.5, ls='--')  # x-axis
plt.axvline(0, color='black',linewidth=0.5, ls='--')  # y-axis
plt.grid(color = 'gray', linestyle = '--', linewidth = 0.5)
plt.legend()
plt.xlim(-10, 10)
plt.ylim(-20, 25)
plt.show()
================================= Tool Message =================================
Name: execute_python


================================== Ai Message ==================================

그래프가 성공적으로 그려졌습니다! 함수 \( y = 2x + 3 \)의 그래프를 확인해 보세요. 이 그래프는 x 값에 따라 y 값이 어떻게 변화하는지를 보여줍니다.

 

 

 

 

다음으로 seoul의 현재 날씨를 물어보며 get_weather tool을 잘 호출하는지도 확인해 보자.

input = {"messages": [("user", "seoul의 현재 날씨는 어떠한가요?")]}
config = {"configurable": {"thread_id": "b"}}
events = graph.stream(input, config, stream_mode="values")

for event in events:
    event["messages"][-1].pretty_print()

================================ Human Message =================================

seoul의 현재 날씨는 어떠한가요?
================================== Ai Message ==================================
Tool Calls:
  get_weather (call_0pZ5Ft9AqcpTx5wQMulQJG90)
 Call ID: call_0pZ5Ft9AqcpTx5wQMulQJG90
  Args:
    location: Seoul
================================= Tool Message =================================
Name: get_weather

The temperature is 5 degrees and it's cloudy.
================================== Ai Message ==================================

서울의 현재 날씨는 기온 5도이며, 흐림입니다.

 

마지막으로, 현재 memory에 해당 thread의 대화 이력이 어떻게 상태에 저장되어 있는지 확인해 보면 앞서 delete_node를 통과하므로 의도한 대로 최근 5개의 메세지가 남아있는 것을 확인할 수 있다.

snapshot = graph.get_state(config)
snapshot.values["messages"]

[
AIMessage(content='그래프가 성공적으로 그려졌습니다! 함수 \\( y = 2x + 3 \\)의 그래프를 확인해 보세요. 이 그래프는 x 값에 따라 y 값이 어떻게 변화하는지를 보여줍니다.', ...),
HumanMessage(content='seoul의 현재 날씨는 어떠한가요?', ...),
ToolMessage(content="The temperature is 5 degrees and it's cloudy.", name='get_weather', ...),
AIMessage(content='서울의 현재 날씨는 기온 5도이며, 흐림입니다.',...)
]

 

 

 

---

출처:

LangGraph. "How to call tools using ToolNode". https://langchain-ai.github.io/langgraph/how-tos/tool-calling/

위키독스 - <랭체인LangChain 노트> - LangChain 한국어 튜토리얼🇰🇷  (https://wikidocs.net/book/14314)