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

[thoffmann-artidis]

I defined a tool_node function that writes variables to the graph state like this:

from typing import Any

from langchain_core.runnables import RunnableConfig
from langchain_core.messages import ToolMessage
from langchain_core.tools import BaseTool
from langgraph.prebuilt.tool_node import _get_state_args

tools = [...]  # put your tools here
tools_by_name: dict[str, BaseTool] = {_tool.name: _tool for _tool in tools}

def tool_node(state: dict, config: RunnableConfig) -> dict[str, Any]:
    """Initialise a tool node that executes the tools and writes
    all returned keys and values to the state variable.

    TODO: This should be a modular class.
    """
    messages = []
    out = {}
    for tool_call in state['messages'][-1].tool_calls:
        tool = tools_by_name[tool_call['name']]
        state_args = {var: state for var in _get_state_args(tool)}
        observation = tool.invoke({**tool_call['args'], **state_args},
                                    config=config)

        message_content = observation.get('output')
        if observation.get('error') is not None:
            message_content = observation.get('error')

        messages.append(ToolMessage(content=message_content,
                                    tool_call_id=tool_call['id']))
        out = {**out, **observation}

    return {'messages': messages, **out}

you can then add it to your langgraph workflow as a node:

workflow.add_node('tools', tool_node)

I'm pretty sure this would not withstand a type checker though.

[makoit]

This is how I have implemented it:

class OperatorAgentGraph:
    def __init__(self, llm: ChatModel, tools: list[StructuredTool]):
        logger.info("Creating SimpleTaskExecutionGraph")
        self.tools = tools
        self.llm = llm
        self.model_with_tools = llm.bind_tools(self.tools, parallel_tool_calls=False)
        self.execution_graph = self.__create_graph()

    def __create_graph(self):
        def should_continue(state: TaskOperatorState) -> TaskOperatorState:
            logger.info(f"Current state: {state}")
            messages = state["messages"]
            last_message = messages[-1]
            if isinstance(last_message, AIMessage) and last_message.tool_calls:
                return "tools"
            return END

        def call_model(state: TaskOperatorState) -> TaskOperatorState:
            # logger.info(f"Current state: {state}")
            messages = state["messages"]
            response = self.model_with_tools.invoke(messages)
            state["messages"] = response
            return state

        def call_tools(state: TaskOperatorState) -> Dict[str, Any]:
            tools_by_name = {tool.name: tool for tool in self.tools}
            result = []
            artifact = {}

            logger.info(f"Current state: {state}")

            # Ensure the last message is an AIMessage with tool_calls
            if (
                not isinstance(state["messages"][-1], AIMessage)
                or not state["messages"][-1].tool_calls
            ):
                raise ValueError(
                    "The last message must be an AIMessage with tool_calls."
                )

            for tool_call in state["messages"][-1].tool_calls:
                tool = tools_by_name[tool_call["name"]]
                content, tool_artifact = tool.invoke(
                    {**tool_call["args"], "state": state}
                )
                result.append(
                    ToolMessage(content=content, tool_call_id=tool_call["id"])
                )

                # Update the artifact dictionary with the tool's artifact
                artifact.update(tool_artifact)

            # Update the state with the new messages and artifacts
            state["messages"].extend(result)
            state.update(artifact)

            return state

        graph = StateGraph(TaskOperatorState)
        graph.add_node("agent", call_model)
        graph.add_node("tools", call_tools)
        graph.add_edge(START, "agent")
        graph.add_conditional_edges("agent", should_continue, ["tools", END])
        graph.add_edge("tools", "agent")
        return graph.compile()

    def invoke(self, initial_state: TaskOperatorState):
        return self.execution_graph.invoke(initial_state)

    async def ainvoke(self, initial_state: TaskOperatorState):
        return await self.execution_graph.ainvoke(initial_state)

And here a sample how a tool looks like:

@tool
def capture_image(
    state: Annotated[TaskOperatorState, InjectedState],
) -> Tuple[str, Dict[str, Any]]:
    """
    Tool explanation
    """

    try:
       # here comes your code....
        
        return (
            "STATUS MESSAGE TEXT",
            {
            "key": "value",
            .....
            },
        )
    except Error as e:
        error_state = f"error: {e}"
        logger.error(error_state)
        return error_state, {}

My tool returns a status message which is passed back to the agent and also a artifact dict which updates the state.

[vbarda]

Hi folks! We have now added support for updating the state from the tools using the new Command type. Please see this how-to guide for reference https://langchain-ai.github.io/langgraph/how-tos/update-state-from-tools/. Let me know if you have any questions / feedback

[suryaashish-tyke]

How to use the state inside a tool_call