How to Implement AI Agents with LangGraph

Introduction

The other day, I wrote an article about how to implement AI agents using a tool called Dify.
https://zenn.dev/asap/articles/a35254020fbbae

Dify makes it incredibly easy to create AI agents and workflows using LLMs, and it's also very simple to embed them into existing web apps or execute them via API.

However, due to my "religious reasons" regarding not being able to use WebUIs (lol), I'm going to try implementing the same thing in Python.

This is a completely self-gratifying article where I implement something with LangGraph just because I like writing code, even though it could just as easily be made with Dify.

If you're using this for professional work or personal projects, Dify is probably the better choice.
From the perspective of flexibility, LangGraph is superior since you can write code directly into the system, but I think that gap will likely shrink soon.
(Especially once Dify supports WebSocket communication.)

However, personally, "just using a tool" is "not fun."
I want to understand the inner workings (which is why I study math and theory in this era of AI tools) and I want to feel like I'm building it myself. So for a hobby, I want to keep using LangGraph.
(Of course, the story would be different if I were doing it to make money.)

LangGraph

LangGraph is a tool built on top of LangChain, and like Dify, it's a framework that allows you to build complex workflows and AI agents using LLMs.

Basically, anything you can do in Dify can be done here, and it can be used with the same syntax regardless of the LLM model.
Also, while it is a feature of LangChain, you don't necessarily have to use LangChain's abstractions to operate the LLM when building AI agents; you can use the native syntax of each LLM provider.

Furthermore, compared to Dify, it is a lower-level framework, making it possible to describe application flows and states with a very high degree of freedom.

On the flip side, the learning cost is higher than Dify.
However, even so, it is far easier than trying to achieve AI agents purely via code without LangGraph. Anyone who can write a bit of code should be able to start using it quickly.

By the way, a specialized book on LangGraph is scheduled to be published on November 9th! I'm looking forward to it.
Introduction to RAG and AI Agents with LangChain and LangGraph [Practical] (Engineer's Selection)

Prerequisites

Environment used

In this article, we will use the Google Colab environment.
Since all LLM processing uses the Azure OpenAI API, there's no need for GPU processing, so we can use the Google Colab environment almost without limits.
(Which is great.)

Generative AI to be Used

In this article, we will use the Azure OpenAI API with gpt-4o as an example.
Of course, you can also use other models (such as gpt-4o-mini), other generative AI providers (such as Anthropic's Claude or Google's Gemini), or local LLMs (such as LLaMa), but in those cases, you will need to modify the model definition section accordingly.

AI Agent to be Created

To implement the same thing as what I created with Dify last time, please refer to the previous article for details, but I will provide a brief explanation.

The AI agent we are creating this time will answer user questions regarding "weather" or "date" within the range of "today, tomorrow, and the day after tomorrow."

More specifically, we will first create the following three workflows:

• tool01
  • The initial LLM determines whether the user's question is about "weather" or "date" and performs conditional branching. Then, the subsequent "LLM specialized in answering today's weather" or "LLM specialized in answering today's date" will provide an answer based on the SYSTEM prompt information (dummy date and dummy weather for today).
• tool02
  • Almost the same as tool01, but answers weather and dates for "tomorrow."
• tool03
  • Almost the same as tool01, but answers weather and dates for "the day after tomorrow."

Next, we will build an AI agent that bundles the above workflows.
The agent triggers the "tool01" workflow if the user's question is "about today," "tool02" if it's "about tomorrow," and "tool03" if it's "about the day after tomorrow." It feeds the user's question into the triggered workflow, retrieves the result, and then provides that answer to the user.

The image looks like this:

(Each tool becomes a workflow)

(Supplement)
If it's just a matter of answering the date or weather, the problem is simple enough that there's little point in using an AI agent.

However, in this workflow, instead of answering the real weather or date, the agent will answer dummy dates and weather based on the information provided in the SYSTEM prompt.
Therefore, for information that the LLM does not inherently know, we can create an AI agent that selects the appropriate workflow, chooses a specialized LLM that has the information within its SYSTEM prompt, and generates the final answer.
(Of course, it's also possible to introduce RAG or similar techniques. RAG is essentially the same in that it adds information to the SYSTEM prompt.)

Does the significance of what we're trying to achieve seem a bit clearer now?

Deliverables

Since this is a standard Google Colab ipynb, I am making the code public as is.
Please see the repository below.
https://github.com/personabb/colab_AI_sample/tree/main/colab_LangGraph_sample

Also, you can see how the AI agent works by running Google Colab as shown below.

https://youtu.be/9_v-R0hxh_Y

The code in this article focuses on LangGraph, so it uses LangChain in a standalone manner when utilizing an LLM within each node.

However, for business or application implementations, syntax like LangChain's LCEL is common, so I recommend using that.

However, I did not adopt it for this article because writing code with LCEL would involve the learning cost of LCEL itself.
For those who want to learn more about LCEL, the following book is an excellent resource.
Introduction to RAG and AI Agents with LangChain and LangGraph [Practical]

I am also writing a sequel to this article. (I will post the URL once it is published.)
How to create an AI agent with LangGraph that asks follow-up questions to the user when input information is insufficient

In the article mentioned above, I have re-implemented the content of this article using LCEL, so please refer to it if you want to know how to write LangGraph using LCEL.

How to Run

Preparation

Preparation for calling AI

Since we are using the Azure OpenAI API this time, the following information is required:

• Endpoint
• API KEY
• API VERSION

Please prepare this information. There is a section to enter it in the notebook.

Store the notebook in your Drive

https://github.com/personabb/colab_AI_sample/blob/main/colab_LangGraph_sample/colab_LangGraph_sample.ipynb

Download (or clone) the "colab_LangGraph_sample.ipynb" mentioned above and upload it to any location in your Google Drive.

Then, launch it as Google Colaboratory. It's OK if the screen looks like this:

Enter the acquired information into the notebook

In the second cell, enter the "OPENAI_API_VERSION", "AZURE_OPENAI_ENDPOINT", and "AZURE_OPENAI_API_KEY" that you acquired during the preparation in the following section:

Also, if you want to use a different Azure model (for example, gpt-4o-mini), change the following part:

model = AzureChatOpenAI(
    azure_deployment="gpt-4o-mini",  # or your deployment
    temperature=0,
)

If you want to use an LLM other than Azure, you will need to significantly modify the following part:

from langchain_openai import AzureChatOpenAI
・・・

os.environ["OPENAI_API_VERSION"] = "2024-08-01-preview"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://xxxxxxxxxx.openai.azure.com"
os.environ["AZURE_OPENAI_API_KEY"] = "xxxxxxxxx"

model = AzureChatOpenAI(
    azure_deployment="gpt-4o",  # or your deployment
    temperature=0,
)

Please modify the above part according to the LangChain syntax of the LLM you want to use.
Conversely, if you set this part correctly, the subsequent processes should work properly regardless of which LLM you use (though performance may vary).

Execution Method

Execute by selecting "Runtime" -> "Run all".

Then, in the very last cell, you will be able to accept user input.

Try asking questions here, just as you did with the AI agent created in Dify.
For example, here is an example of a question I asked:

(Text)

Your message: 明後日は傘が必要ですか？
Assistant: 明後日は傘が必要ですか？

明後日は晴れですので、傘は必要ありません。
Your message: 明日は緑の日ですか？
Assistant: 明日は緑の日ですか？

いいえ、明日は緑の日ではありません。緑の日は日本では5月4日に祝われます。明日は10月24日です。
Your message: 明日は傘を持って行かなくてもいいですか？
Assistant: 明日は傘を持って行かなくてもいいですか？

明日は曇りですので、傘を持って行かなくても大丈夫かもしれませんが、念のため天気予報を確認しておくと安心です。
Your message: exit
Goodbye!

As you can see, by selecting the appropriate workflow and calling the appropriate expert LLM, the agent is answering using information that only that LLM knows.
(For example, the information that it will be sunny the day after tomorrow, and the weather/date for tomorrow, is dummy data I embedded in the prompt, so it's information known only to the specific (single) LLM.)

Finally, you can also terminate the system by typing "exit".

Code Explanation

Now that we have seen that AI agents can be implemented using LangGraph, I will explain how it is built.

Basically, I have also included comments within the code, so I will provide a brief explanation here.

Cell 1 (Installing Packages)

colab_LangGraph_sample.ipynb

!pip install langgraph langchain_openai langchain_core httpx

We are installing the necessary packages.
These are the primary packages required to use LangGraph and LangChain.

Cell 2 (Importing Modules and Model Definition)

colab_LangGraph_sample.ipynb

import os
import glob
from typing import Annotated, List, Literal
from langgraph.graph import StateGraph, END
from langchain_openai import AzureChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from typing_extensions import TypedDict ,Optional, List, Dict
from langchain_core.pydantic_v1 import BaseModel, Field

os.environ["OPENAI_API_VERSION"] = "2024-08-01-preview"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://xxxxxxxxxx.openai.azure.com"
os.environ["AZURE_OPENAI_API_KEY"] = "xxxxxxxxx"


model = AzureChatOpenAI(
    azure_deployment="gpt-4o",  # or your deployment
    temperature=0,
)

Here, we define the gpt-4o model that was previously deployed on Azure using the Azure OpenAI API.
For this, please enter the AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT that you obtained.
(Of course, you can also set them as environment variables. In fact, it is better to do so.)

In Colab, you can register environment variables from the key icon in the left tab and call them using the following command:

from google.colab import userdata
os.environ["OPENAI_API_VERSION"] = userdata.get('OPENAI_API_VERSION')
os.environ["AZURE_OPENAI_ENDPOINT"] = userdata.get('AZURE_OPENAI_ENDPOINT')
os.environ["AZURE_OPENAI_API_KEY"] = userdata.get('AZURE_OPENAI_API_KEY')

By the way, if you want to use a more cost-effective model like gpt-4o-mini, you can do so by changing it as follows:
(Assuming, of course, that it is deployed on the Azure side.)

model = AzureChatOpenAI(
    azure_deployment="gpt-4o-mini",  # or your deployment
    temperature=0,
)

Cell 3 (Specifying Data Types)

colab_LangGraph_sample.ipynb

# Specify the type of data flowing through the workflow.
class State(TypedDict, total=False):
    message_type: Optional[str]
    messages: Optional[str]

# Specify the class to fix the output of the classifier LLM
class MessageType(BaseModel):
    message_type: str = Field(description="The type of the message", example="search")

# Specify the class to fix the selection output of the tools chosen by the AI agent
class ToolType(BaseModel):
    message_type: str = Field(description="Must use the Tool", example="tool01")

# Specify the type of data flowing through the workflow.
class State(TypedDict, total=False):
    message_type: Optional[str]
    messages: Optional[str]

The State class specifies the type of data that flows throughout the entire workflow.
We proceed with the workflow by storing the selected tool name in message_type and the user's question content in messages.

# Specify the class to fix the output of the classifier LLM
class MessageType(BaseModel):
    message_type: str = Field(description="The type of the message", example="search")

# Specify the class to fix the selection output of the tools chosen by the AI agent
class ToolType(BaseModel):
    message_type: str = Field(description="Must use the Tool", example="tool01")

The above classes are for structuring the values output by the LLM.
The model will produce structured output while also referring to the explanations provided in description.

Cell 4 (Defining Agent LLM Behavior)

colab_LangGraph_sample.ipynb

# Define a function for the AI agent to classify the user's question (today, tomorrow, or the day after tomorrow) and call the appropriate workflow (tool01 to tool03)

# To fix the output
tools = model.with_structured_output(ToolType)

def select_tool(State):
    # Create prompt
    classification_prompt = """
    ## You are a message classifier.
    ## For questions about today, reply with "tool01".
    ## For questions about tomorrow, reply with "tool02".
    ## For questions about the day after tomorrow, reply with "tool03".

    ## user message: {user_message}
    """

    if State["messages"]:
        return {
            "message_type": tools.invoke(classification_prompt.format(user_message=State["messages"])).message_type,
            "messages": State["messages"]
            }
    else:
        return {"message": "No user input provided"}

Here, the initial AI agent determines whether the user's question is "about today," "about tomorrow," or "about the day after tomorrow" and asks which workflow (tool) to select.

Also, this LLM is made to produce structured output using the following:

tools = model.with_structured_output(ToolType)

In the return statement, the output basically follows the State class presented above.
Since the first LLM outputs which tool to select based on the user's question, the model's output is assigned to message_type, and messages contains the user's question content as is.

return {
    "message_type": tools.invoke(classification_prompt.format(user_message=State["messages"])).message_type,
    "messages": State["messages"]
    }

Cell 5 (Defining the Date/Weather Question Classification LLM within the Workflow)

colab_LangGraph_sample.ipynb

# Define the LLM node to classify user questions (weather or date), used in tool01 to tool03

# To fix the output
classifier = model.with_structured_output(MessageType)

# Define the function to be executed in the node
def classify(State):
    # Create prompt
    classification_prompt = """
    ## You are a message classifier.
    ## If the user is asking about the weather, reply with "weather".
    ## For any other questions, reply with "day".

    ## user message: {user_message}
    """

    if State["messages"]:
        return {
            "message_type": classifier.invoke(classification_prompt.format(user_message=State["messages"])).message_type,
            "messages": State["messages"]
            }
    else:
        return {"message": "No user input provided"}

The classifier itself only needs to distinguish between "questions about weather" and "questions about date" regardless of the date, so it is used commonly.

Similar to Cell 4, the output is structured using with_structured_output.

The behavior is defined by the classify function.

Cell 6 (Defining Expert LLMs within the Workflow)

colab_LangGraph_sample.ipynb

# Define the question-answering LLM node for the tool01 workflow
def chat_w1(State):
    if State["messages"]:
        State["messages"][0] = ("system", "Please repeat the user's question. Then, answer the question. Note that it is raining today.")
        return {"messages": model.invoke(State["messages"])}
    return {"messages": "No user input provided"}

def chat_d1(State):
    if State["messages"]:
        State["messages"][0] = ("system", "After repeating the user's question, please answer it. Note that today is 10/23.")
        return {"messages": model.invoke(State["messages"])}
    return {"messages": "No user input provided"}

# Define the question-answering LLM node for the tool02 workflow
def chat_w2(State):
    if State["messages"]:
        State["messages"][0] = ("system", "Please repeat the user's question. Then, answer the question. Note that it is cloudy tomorrow.")
        return {"messages": model.invoke(State["messages"])}
    return {"messages": "No user input provided"}

def chat_d2(State):
    if State["messages"]:
        State["messages"][0] = ("system", "After repeating the user's question, please answer it. Note that tomorrow is 10/24.")
        return {"messages": model.invoke(State["messages"])}
    return {"messages": "No user input provided"}

# Define the question-answering LLM node for the tool03 workflow
def chat_w3(State):
    if State["messages"]:
        State["messages"][0] = ("system", "Please repeat the user's question. Then, answer the question. Note that it will be sunny the day after tomorrow.")
        return {"messages": model.invoke(State["messages"])}
    return {"messages": "No user input provided"}

def chat_d3(State):
    if State["messages"]:
        State["messages"][0] = ("system", "After repeating the user's question, please answer it. Note that the day after tomorrow is 10/25.")
        return {"messages": model.invoke(State["messages"])}
    return {"messages": "No user input provided"}

# Define the final output node
def response(State):
    return State

Here, after the question is classified as either weather or date, a specialized LLM answers based on the expert knowledge described in the SYSTEM prompt.

Additionally, a node for the final output is defined at the end.

Cell 7 (Defining the Flow of the Three Workflows)

colab_LangGraph_sample.ipynb

# Define the type of data flowing through the sub-nodes (child nodes) for organization.

# Define the type of data flowing through child nodes
class ChildState(TypedDict, total=False):
    message_type: Optional[str]
    messages: Optional[str]

# Define nodes and edges for the tool01 workflow
# Adding nodes
child_builder1 = StateGraph(ChildState)
child_builder1.add_node("classify1", classify)
child_builder1.add_node("chat_w1", chat_w1)
child_builder1.add_node("chat_d1", chat_d1)
child_builder1.add_node("response1", response)

# Adding edges
child_builder1.add_edge("chat_d1", "response1")
child_builder1.add_edge("chat_w1", "response1")
# Conditional branching
child_builder1.add_conditional_edges("classify1", lambda state: state["message_type"], {"weather": "chat_w1", "day": "chat_d1"})

# Specifying entry and finish points
child_builder1.set_entry_point("classify1")
child_builder1.set_finish_point("response1")

# Define nodes and edges for the tool02 workflow
# Adding nodes
child_builder2 = StateGraph(ChildState)
child_builder2.add_node("classify2", classify)
child_builder2.add_node("chat_w2", chat_w2)
child_builder2.add_node("chat_d2", chat_d2)
child_builder2.add_node("response2", response)

# Adding edges
child_builder2.add_edge("chat_d2", "response2")
child_builder2.add_edge("chat_w2", "response2")
# Conditional branching
child_builder2.add_conditional_edges("classify2", lambda state: state["message_type"], {"weather": "chat_w2", "day": "chat_d2"})

# Specifying entry and finish points
child_builder2.set_entry_point("classify2")
child_builder2.set_finish_point("response2")

# Define nodes and edges for the tool03 workflow
# Adding nodes
child_builder3 = StateGraph(ChildState)
child_builder3.add_node("classify3", classify)
child_builder3.add_node("chat_w3", chat_w3)
child_builder3.add_node("chat_d3", chat_d3)
child_builder3.add_node("response3", response)

# Adding edges
child_builder3.add_edge("chat_d3", "response3")
child_builder3.add_edge("chat_w3", "response3")
# Conditional branching
child_builder3.add_conditional_edges("classify3", lambda state: state["message_type"], {"weather": "chat_w3", "day": "chat_d3"})

# Specifying entry and finish points
child_builder3.set_entry_point("classify3")
child_builder3.set_finish_point("response3")

In this section, we define the flow for each of the three workflows. Since the same process is repeated three times, I will focus on explaining one of them.

# Define the type of data flowing through child nodes
class ChildState(TypedDict, total=False):
    message_type: Optional[str]
    messages: Optional[str]

Here, we define the type of data that flows within the workflow. Basically, it is the same as the parent agent LLM's state.

child_builder1 = StateGraph(ChildState)

This is where the workflow is created. A workflow named child_builder1 is initialized. The data structure defined in ChildState will flow through this workflow.

child_builder1.add_node("classify1", classify)
child_builder1.add_node("chat_w1", chat_w1)
child_builder1.add_node("chat_d1", chat_d1)
child_builder1.add_node("response1", response)

Here, we add "nodes" to the workflow graph. A node executes the processing logic of its assigned function.

In the example above:

• classify1 node: Processes the classification of the question content.
• chat_w1 node: An expert LLM processes weather-related questions.
• chat_d1 node: An expert LLM processes date-related questions.
• response1 node: The final output node.

The input and output of these nodes follow the ChildState class definition.

# Adding edges
child_builder1.add_edge("chat_d1", "response1")
child_builder1.add_edge("chat_w1", "response1")

Next, we define the connections between nodes as "edges".
In the code above:

• An edge connects chat_d1 to response1.
• An edge connects chat_w1 to response1.
  These connections are straightforward and fixed without any logic, allowing for this simple definition.

On the other hand, for conditional logic, we define it as follows:

child_builder1.add_conditional_edges("classify1", lambda state: state["message_type"], {"weather": "chat_w1", "day": "chat_d1"})

This allows us to define which node the data flows to from classify1 based on the value of state["message_type"]. If the question is about weather, the chat_w1 node is called; if it's about the date, chat_d1 is called.

# Specifying entry and finish points
child_builder1.set_entry_point("classify1")
child_builder1.set_finish_point("response1")

The start and end nodes of the graph are defined here. Through these steps, we have connected the edges in order from the initial node to the final one. This is how a graph is defined. The second and third workflows are defined exactly the same way.

Cell 8 (Defining the Agent's Flow)

colab_LangGraph_sample.ipynb

# Define the StateGraph of the AI agent using the workflows defined as child nodes

# Adding nodes
graph_builder = StateGraph(State)
graph_builder.add_node("select_tool", select_tool)
graph_builder.add_node("start_tool01", child_builder1.compile())
graph_builder.add_node("start_tool02", child_builder2.compile())
graph_builder.add_node("start_tool03", child_builder3.compile())
graph_builder.add_node("response", response)

# Adding edges
graph_builder.add_edge("start_tool01", "response")
graph_builder.add_edge("start_tool02", "response")
graph_builder.add_edge("start_tool03", "response")
# Conditional branching
graph_builder.add_conditional_edges("select_tool", lambda state: state["message_type"], {"tool01": "start_tool01", "tool02": "start_tool02", "tool03": "start_tool03"})

# Specifying entry and finish points
graph_builder.set_entry_point("select_tool")
graph_builder.set_finish_point("response")

# Build graph
graph = graph_builder.compile()

Next, we define the graph for the agent LLM that bundles the three workflows together. Basically, we define the nodes and edges just as explained in the previous cell.

Since the individual workflows were already defined as graphs in the previous cell, we can define those graphs themselves as nodes here. It is written like this:

graph_builder.add_node("start_tool01", child_builder1.compile())

When calling a workflow that has been made into a graph, you specify the graph name (e.g., child_builder1) and use .compile() to create an instance of the CompiledGraph class to be treated as a node. Finally, once all graphs are defined, you build the overall graph using the .compile() method.

# Build graph
graph = graph_builder.compile()

Cell 9 (Graph Visualization)

colab_LangGraph_sample.ipynb

# Visualize the final graph
from IPython.display import Image, display

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

The graph can be visualized in the manner shown above.
The visualization for this instance is as follows:

Cell 10 (Executing the Conversation)

colab_LangGraph_sample.ipynb

# Interaction loop
user_input = input("Your message: ")
while user_input.lower() != "exit":

        # Align input to the State structure
        state = {
            "messages": [
                    ("system",""),

                    ("human", user_input),
            ]
        }

        # Retrieve and display response
        last_content = None

        for event in graph.stream(state):
            # Get the content of 'messages' from the last 'response'
            if "response" in event and "messages" in event["response"]:
                messages = event["response"]["messages"]
                # If it is an AIMessage object, get the content attribute
                if hasattr(messages, "content"):
                    last_content = messages.content

        # Display the final message
        if last_content:
            print("Assistant:", last_content)
        else:
            print("No content found.")

        user_input = input("Your message: ")
print("Goodbye!")

In the code above, while the console accepts user input, the entered question is set into the State format and fed into the graph constructed via graph.stream(state). In the State class this time, message_type is optional, so it is fine if it is not present initially. Therefore, at the start, only the user's input is fed into the graph.

Also, looking at the state variable, an empty system prompt is prepared in addition to the user input. As you can see by looking at the functions where the LLM performs processing, the following type of logic actually intervenes before a message is input into the model:

State["messages"][0] = ("system", "Please repeat the user's question. Then, answer the question. However, it will be sunny the day after tomorrow.")

In other words, an empty element is provided so that each LLM can overwrite it with its own system prompt.

Finally, from the values returned from the graph, only the LLM's response text part is extracted and displayed on the console.

Additionally, the system is designed to allow for repeated processing until the user enters "exit".

Summary

Thank you very much for reading this far!

As expected, I felt that using LangGraph to do the same thing as Dify is quite a bit of work.
However, because you can write it out as code, the internal degree of freedom should be high.
With Dify, you can only interact with the system via API calls, but with LangGraph, you can write directly into the system or interact through connection methods not supported by Dify, such as WebSocket communication.

I hope that future LangGraph updates will make it even easier to use!

References

https://zenn.dev/pharmax/articles/148bc9497d68dd