Tool calling in LangGraph and how to update the Graph-State with Tools output

[dominik-roemer]

Hey,
atm I'm looking into LangGraph and the agents ability to call tools. I was able to build an agent which has this architecture. The idea is taken from: https://langchain-ai.github.io/langgraph/tutorials/introduction/#part-2-enhancing-the-chatbot-with-tools

When a tool is called the Chatbot receives a Toolmesage with the output of the tool. Now my question is, how to take this output and write it into the state. Let's say my state looks something like this:

messages: Annotated[list[AnyMessage], add_messages]
   first_name: str
   last_name: str
   date_of_birth: str

And I have a very simple tool which is called when the user provides his name in the query, to find the date_of_birth in the database.

find_user_date_of_birth(user_name: str):
"""This tool is used to search for the users date of birth in the customer DB if the user name is given in the query."""
  date_of_birth = "01.01.1990" # this is just a mock tool
  return f"The user {username} birthday is {date_of_birth}"

How am I now able to update the GraphState in a way that date_of_birth is set in the State? The tool node is returning {"messages": output} which will update the state. But what if I want to return "date_of_birth"? Is there an easy way to solve this dynamically or is there even a better way?

[isahers1]

The prebuilt ToolNode will only append to the messages key in your state, but you can define a custom node like so:

def custom_tool_node(state):
    return {"date_of_birth": find_user_date_of_birth(state['first_name']+state['last_name'])}

This node will update the "date_of_birth" key. You may also wish to append the ToolMessage to the list of messages, so be careful of that as well. Let me know if this helps!

[thoffmann-artidis]

Is there any reason why ToolNode doesn't write to any other state variable than messages? It's a widespread use case where you would like to retrieve some data through a tool call and pass it to another agent.

[makoit]

I have a similar question and opened a question for it here. @isahers1 can you explain more in detail how I need to implement a tool node by my own? I have following code:

import getpass
import os
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import InjectedState
from langchain_core.tools import tool
from typing import Annotated
from langgraph.prebuilt.chat_agent_executor import AgentState
from typing_extensions import Annotated, TypedDict, NotRequired
from typing import Callable, Literal, Optional, Sequence, Type, TypeVar, Union, cast
from langchain_core.messages import AIMessage, BaseMessage, SystemMessage, ToolMessage,  AnyMessage
from langgraph.graph.message import add_messages
from langgraph.managed import IsLastStep, RemainingSteps
from langgraph.graph import MessagesState
from langchain_openai import ChatOpenAI


class MyState(TypedDict):
    image_data: NotRequired[bytes]
    messages: Annotated[list[AnyMessage], add_messages]


@tool
def capture_image(state: Annotated[MyState, InjectedState]) -> str:
    """Capture an image and store it in the state."""
    print("STATE CAPTURE IMAGE: ", state)
    with open('/workspace/picker/test_images/image.jpeg', 'rb') as image_file:
        image = image_file.read()
    # Replace with actual API call to capture image
    return "Image captured and stored in state."


@tool
def analyze_image(state: Annotated[MyState, InjectedState]) -> str:
    """Analyze the image stored in the state."""
    print("STATE ANALZYE IMAGE: ", state)
    if not state["image_data"]:
        return "No image data found in state."
    # Replace with actual API call to analyze image
    analysis_result = "Image analysis result"  # Mock analysis result
    return analysis_result


model = ChatOpenAI(
    model="gpt-4o", temperature=0
)

# Define tools and bind them to model
tools = [capture_image, analyze_image]
model_with_tools = model.bind_tools(tools, parallel_tool_calls=False)

# Define the function that determines whether to continue or not


def should_continue(state: MyState) -> MyState:
    print("STATE SHOULD CONTINUE: ", state)
    messages = state["messages"]
    last_message = messages[-1]
    print("LAST MESSAGE: ", last_message)
    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        return "tools"
    return END

# Define the function that calls the model

def call_model(state: MyState) -> MyState:
    print("STATE IN CALL MODEL: ", state)
    messages = state["messages"]
    response = model_with_tools.invoke(messages)
    state["messages"] = response
    return state

# Define the graph
workflow = StateGraph(MyState)
workflow.add_node("agent", call_model)
tool_node = ToolNode(tools)
workflow.add_node("tools", tool_node)
workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", should_continue, ["tools", END])
workflow.add_edge("tools", "agent")

# Compile the graph
app = workflow.compile()


# Initialize the state
initial_state = {
    "messages": [HumanMessage(content="Please capture an image and explain me what you can see.")]
}

# Execute the graph
result = app.invoke(initial_state)

print(result)

I want to share data from tools in the graph state. so that one tool can capture a image and other one can consume and analyze the image. For this a inject the state into the tool, but I cannot update the state in the tool, so they way back is missing.

[jasontjahjono]

+1, waiting for an answer for tools to save to state instead of pass results back to agent