langgraph/examples/react-agent-structured-output.ipynb

How to force function calling agent to structure output

You might want your agent to return its output in a structured format. For example, if the output of the agent is used by some other downstream software, you may want the output to be in the same structured format every time the agent is invoked to ensure consistency.

This guide shows how you can do this. We will be using a basic ReAct agent (a model node and a tool-calling node) together with a third node at the end that will format response for the user.

Setup

Structured Output

First we need to define how we want to structure our output. To do this, we will use the with_structured_output method from LangChain, which you can read more about here.

In [1]:

from pydantic import BaseModel, Field

class WeatherResponse(BaseModel):
    temperature: float = Field(description="The temperature in fahrenheit")
    wind_directon: str = Field(description="The direction of the wind in abbreviated form")
    wind_speed: float = Field(description="The speed of the wind in km/h")

Graph State

We can now define our graph state:

In [2]:

from typing import Annotated, Any
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    # list of chat messages from user, LLM, and tools
    messages: Annotated[list, add_messages]
    # Final structured response from the agent
    final_response: WeatherResponse

Tools + Models

We can now instantiate the tools and models we are going to use in our graph. We are going to use a single tool in this example for finding the weather, and we are going to have two models in our graph, one that does the function calling and one that does the responding.

In [3]:

from typing import Literal
from langchain_core.tools import tool
from langchain_anthropic import ChatAnthropic

@tool
def get_weather(city: Literal["nyc", "sf"]):
    """Use this to get weather information."""
    if city == "nyc":
        return "It is cloudy in NYC, with 5 mph winds in the North-East direction and a temperature of 70 degrees"
    elif city == "sf":
        return "It is 75 degrees and sunny in SF, with 3 mph winds in the South-East direction"
    else:
        raise AssertionError("Unknown city")
    
tools = [get_weather]
    
model = ChatAnthropic(model="claude-3-opus-20240229")
    
model_with_tools = model.bind_tools(tools)
model_with_structured_output = model.with_structured_output(WeatherResponse)

Define Graph

Now that we have defined our tools and models, we can define our graph.

In [18]:

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_core.messages import HumanMessage

# Define the function that determines whether to continue or not
def should_continue(state: AgentState):
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no function call, then we respond to the user
    if not last_message.tool_calls:
        return "respond"
    # Otherwise if there is, we continue
    else:
        return "continue"

# Define the function that calls the model
def call_model(state: AgentState):
    response = model_with_tools.invoke(state['messages'])
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}

# Define the function that responds to the user
def respond(state: AgentState):
    # We call the model with structured output in order to return the same format to the user every time
    # state['messages'][-2] is the last ToolMessage in the convo, which we convert to a HumanMessage for the model to use
    # We could also pass the entire chat history, but this saves tokens since all we care to structure is the output of the tool
    response = model_with_structured_output.invoke([HumanMessage(content=state['messages'][-2].content)])
    # We return the final answer
    return {"final_response": response}


# Define a new graph
workflow = StateGraph(AgentState)

# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("respond", respond)
workflow.add_node("tools", ToolNode(tools))

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")

# We now add a conditional edge
workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "continue": "tools",
        "respond": "respond",
    },
)

workflow.add_edge("tools", "agent")
workflow.add_edge("respond", END)
graph = workflow.compile()

Usage

We can now invoke our graph to verify that the output is being structured as desired:

In [19]:

answer = graph.invoke(input={"messages": [("human", "what's the weather in SF?")]})['final_response']

In [ ]:

answer

WeatherResponse(temperature=75.0, wind_directon='SE', wind_speed=3.0)

As we can see, the agent returned a WeatherResponse object as we expected. If would now be easy to use this agent in a more complex software stack without having to worry about the output of the agent not matching the format expected from the next step in the stack.