Creating Task-Oriented Dialog systems with LangGraph and LangChain

A Task-Oriented Dialogue system (ToD) is a system that assists users in achieving a selected task, comparable to booking a restaurant, planning a travel itinerary or ordering delivery food.

We all know that we instruct LLMs using prompts, but how can we implement these ToD systems in order that the conversation all the time revolves across the task we would like the users to realize? A method of doing that’s through the use of prompts, memory and tool calling. FortunatelyLangChain + LangGraph may also help us tie all these items together.

In this text, you’ll learn tips on how to construct a Task Oriented Dialogue System that helps users create User Stories with a high level of quality. The system is all based on LangGraph’s Prompt Generation from User Requirements tutorial.

On this tutorial we assume you already know tips on how to use LangChain. A User Story has some components like objective, success criteria, plan of execution and deliverables. The user should provide each of them, and we’d like to “hold their hand” into providing them one after the other. Doing that using only LangChain would require quite a lot of ifs and elses.

With LangGraph we are able to use a graph abstraction to create cycles to manage the dialogue. It also has built-in persistence, so we don’t must worry about actively tracking the interactions that occur throughout the graph.

The foremost LangGraph abstraction is the StateGraph, which is used to create graph workflows. Each graph must be initialized with a state_schema: a schema class that every node of the graph uses to read and write information.

The flow of our system will consist of rounds of LLM and user messages. The foremost loop will contain these steps:

1. User says something
2. LLM reads the messages of the state and decides if it’s able to create the User Story or if the user should respond again

Our system is easy so the schema consists only of the messages that were exchanged within the dialogue.

from langgraph.graph.message import add_messages
class StateSchema(TypedDict):
messages: Annotated[list, add_messages]

The add_messages method is used to merge the output messages from each node into the present list of messages within the graph’s state.

Speaking about nodes, one other two foremost LangGraph concepts are Nodes and Edges. Each node of the graph runs a function and every edge controls the flow of 1 node to a different. We even have START and END virtual nodes to inform the graph where to begin the execution and where the execution should end.

To run the system we’ll use the .stream() method. After we construct the graph and compile it, each round of interaction will undergo the START until the END of the graph and the trail it takes (which nodes should run or not) is controlled by our workflow combined with the state of the graph. The next code has the foremost flow of our system:

config = {"configurable": {"thread_id": str(uuid.uuid4())}}
while True:
user = input("User (q/Q to quit): ")
if user in {"q", "Q"}:
print("AI: Byebye")
break
output = None
for output in graph.stream(
{"messages": [HumanMessage(content=user)]}, config=config, stream_mode="updates"
):
last_message = next(iter(output.values()))["messages"][-1]
last_message.pretty_print()
if output and "prompt" in output:
print("Done!")

At each interaction (if the user didn’t type “q” or “Q” to quit) we run graph.stream() passing the message of the user using the “updates” stream_mode, which streams the updates of the state after each step of the graph (https://langchain-ai.github.io/langgraph/concepts/low_level/#stream-and-astream). We then get this last message from the state_schema and print it.

On this tutorial we’ll still learn tips on how to create the nodes and edges of the graph, but first let’s talk more concerning the architecture of ToD systems normally and learn tips on how to implement one with LLMs, prompts and tool calling.

The foremost components of a framework to construct End-to-End Task-Oriented Dialogue systems are [1]:

1. Natural Language Understanding (NLU) for extracting the intent and key slots of users
2. Dialogue State Tracking (DST) for tracing users’ belief state given dialogue
3. Dialogue Policy Learning (DPL) to determine the subsequent step to take
4. Natural Language Generation (NLG) for generating dialogue system response

By utilizing LLMs, we are able to mix a few of these components into just one. The NLP and the NLG components are easy peasy to implement using LLMs since understanding and generating dialogue responses are their specialty.

We will implement the Dialogue State Tracking (DST) and the Dialogue Policy Learning (DPL) through the use of LangChain’s SystemMessage to prime the AI behavior and all the time pass this message each time we interact with the LLM. The state of the dialogue must also all the time be passed to the LLM at every interaction with the model. Which means we are going to ensure that the dialogue is all the time centered across the task we would like the user to finish by all the time telling the LLM what the goal of the dialogue is and the way it should behave. We’ll do this first through the use of a prompt:

prompt_system_task = """Your job is to assemble information from the user concerning the User Story they should create.
You must obtain the next information from them:
- Objective: the goal of the user story. must be concrete enough to be developed in 2 weeks.
- Success criteria the sucess criteria of the user story
- Plan_of_execution: the plan of execution of the initiative
- Deliverables: the deliverables of the initiative
In case you aren't capable of discern this info, ask them to make clear! Don't try and wildly guess. 
Each time the user responds to one among the standards, evaluate whether it is detailed enough to be a criterion of a User Story. If not, ask inquiries to help the user higher detail the criterion.
Don't overwhelm the user with too many questions directly; ask for the knowledge you would like in a way that they should not have to write down much in each response. 
At all times remind them that in the event that they have no idea tips on how to answer something, you may also help them.
After you're capable of discern all the knowledge, call the relevant tool."""

After which appending this prompt everytime we send a message to the LLM:

def domain_state_tracker(messages):
return [SystemMessage(content=prompt_system_task)] + messages

One other vital concept of our ToD system LLM implementation is tool calling. In case you read the last sentence of the prompt_system_task again it says “After you’re capable of discern all the knowledge, call the relevant tool”. This fashion, we’re telling the LLM that when it decides that the user provided all of the User Story parameters, it should call the tool to create the User Story. Our tool for that will likely be created using a Pydantic model with the User Story parameters.

By utilizing only the prompt and power calling, we are able to control our ToD system. Beautiful right? Actually we also need to make use of the state of the graph to make all this work. Let’s do it in the subsequent section, where we’ll finally construct the ToD system.

Alright, time to do some coding. First we’ll specify which LLM model we’ll use, then set the prompt and bind the tool to generate the User Story:

import os
from dotenv import load_dotenv, find_dotenv
from langchain_openai import AzureChatOpenAI
from langchain_core.pydantic_v1 import BaseModel
from typing import List, Literal, Annotated
_ = load_dotenv(find_dotenv()) # read local .env file
llm = AzureChatOpenAI(azure_deployment=os.environ.get("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"),
openai_api_version="2023-09-01-preview",
openai_api_type="azure",
openai_api_key=os.environ.get('AZURE_OPENAI_API_KEY'),
azure_endpoint=os.environ.get('AZURE_OPENAI_ENDPOINT'),
temperature=0)
prompt_system_task = """Your job is to assemble information from the user concerning the User Story they should create.
You must obtain the next information from them:
- Objective: the goal of the user story. must be concrete enough to be developed in 2 weeks.
- Success criteria the sucess criteria of the user story
- Plan_of_execution: the plan of execution of the initiative
In case you aren't capable of discern this info, ask them to make clear! Don't try and wildly guess. 
Each time the user responds to one among the standards, evaluate whether it is detailed enough to be a criterion of a User Story. If not, ask inquiries to help the user higher detail the criterion.
Don't overwhelm the user with too many questions directly; ask for the knowledge you would like in a way that they should not have to write down much in each response. 
At all times remind them that in the event that they have no idea tips on how to answer something, you may also help them.
After you're capable of discern all the knowledge, call the relevant tool."""
class UserStoryCriteria(BaseModel):
"""Instructions on tips on how to prompt the LLM."""
objective: str
success_criteria: str
plan_of_execution: str
llm_with_tool = llm.bind_tools([UserStoryCriteria])

As we were talking earlier, the state of our graph consists only of the messages exchanged and a flag to know if the user story was created or not. Let’s create the graph first using StateGraph and this schema:

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
class StateSchema(TypedDict):
messages: Annotated[list, add_messages]
created_user_story: bool
workflow = StateGraph(StateSchema)

The subsequent image shows the structure of the ultimate graph:

At the highest now we have a talk_to_user node. This node can either:

• Finalize the dialogue (go to the finalize_dialogue node)
• Determine that it’s time to attend for the user input (go to the END node)

Because the foremost loop runs endlessly (while True), each time the graph reaches the END node, it waits for the user input again. This may grow to be more clear after we create the loop.

Let’s create the nodes of the graph, starting with the talk_to_user node. This node needs to maintain track of the duty (maintaing the foremost prompt during all of the conversation) and in addition keep the message exchanges since it’s where the state of the dialogue is stored. This state also keeps which parameters of the User Story are already filled or not using the messages. So this node should add the SystemMessage each time and append the brand new message from the LLM:

def domain_state_tracker(messages):
return [SystemMessage(content=prompt_system_task)] + messages
def call_llm(state: StateSchema):
"""
talk_to_user node function, adds the prompt_system_task to the messages,
calls the LLM and returns the response
"""
messages = domain_state_tracker(state["messages"])
response = llm_with_tool.invoke(messages)
return {"messages": [response]}

Now we are able to add the talk_to_user node to this graph. We’ll do this by giving it a reputation after which passing the function we’ve created:

workflow.add_node("talk_to_user", call_llm)

This node must be the primary node to run within the graph, so let’s specify that with an edge:

workflow.add_edge(START, "talk_to_user")

To this point the graph looks like this:

To manage the flow of the graph, we’ll also use the message classes from LangChain. Now we have 4 kinds of messages:

• SystemMessage: message for priming AI behavior
• HumanMessage: message from a human
• AIMessage: the message returned from a chat model as a response to a prompt
• ToolMessage: message containing the results of a tool invocation, used for passing the results of executing a tool back to a model

We’ll use the sort of the last message of the graph state to manage the flow on the talk_to_user node. If the last message is an AIMessage and it has the tool_calls key, then we’ll go to the finalize_dialogue node since it’s time to create the User Story. Otherwise, we must always go to the END node because we’ll restart the loop because it’s time for the user to reply.

The finalize_dialogue node should construct the ToolMessage to pass the result to the model. The tool_call_id field is used to associate the tool call request with the tool call response. Let’s create this node and add it to the graph:

def finalize_dialogue(state: StateSchema):
"""
Add a tool message to the history so the graph can see that it`s time to create the user story
"""
return {
"messages": [
ToolMessage(
content="Prompt generated!",
tool_call_id=state["messages"][-1].tool_calls[0]["id"],
)
]
}
workflow.add_node("finalize_dialogue", finalize_dialogue)

Now let’s create the last node, the create_user_story one. This node will call the LLM using the prompt to create the User Story and the knowledge that was gathered throughout the conversation. If the model decided that it was time to call the tool then the values of the important thing tool_calls must have all the information to create the User Story.

prompt_generate_user_story = """Based on the next requirements, write a superb user story:
{reqs}"""
def build_prompt_to_generate_user_story(messages: list):
tool_call = None
other_msgs = []
for m in messages:
if isinstance(m, AIMessage) and m.tool_calls: #tool_calls is from the OpenAI API
tool_call = m.tool_calls[0]["args"]
elif isinstance(m, ToolMessage):
proceed
elif tool_call isn't None:
other_msgs.append(m)
return [SystemMessage(content=prompt_generate_user_story.format(reqs=tool_call))] + other_msgs
def call_model_to_generate_user_story(state):
messages = build_prompt_to_generate_user_story(state["messages"])
response = llm.invoke(messages)
return {"messages": [response]}
workflow.add_node("create_user_story", call_model_to_generate_user_story)

With all of the nodes are created, it’s time so as to add the edges. We’ll add a conditional edge to the talk_to_user node. Keep in mind that this node can either:

• Finalize the dialogue if it’s time to call the tool (go to the finalize_dialogue node)
• Determine that we’d like to assemble user input (go to the END node)

Which means we’ll only check if the last message is an AIMessage and has the tool_calls key; otherwise we must always go to the END node. Let’s create a function to envision this and add it as an edge:

def define_next_action(state) -> Literal["finalize_dialogue", END]:
messages = state["messages"]
if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:
return "finalize_dialogue"
else:
return END
workflow.add_conditional_edges("talk_to_user", define_next_action)

Now let’s add the opposite edges:

workflow.add_edge("finalize_dialogue", "create_user_story")
workflow.add_edge("create_user_story", END)

With that the graph workflow is finished. Time to compile the graph and create the loop to run it:

memory = MemorySaver()
graph = workflow.compile(checkpointer=memory)
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
while True:
user = input("User (q/Q to quit): ")
if user in {"q", "Q"}:
print("AI: Byebye")
break
output = None
for output in graph.stream(
{"messages": [HumanMessage(content=user)]}, config=config, stream_mode="updates"
):
last_message = next(iter(output.values()))["messages"][-1]
last_message.pretty_print()
if output and "create_user_story" in output:
print("User story created!")

Let’s finally test the system:

With LangGraph and LangChain we are able to construct systems that guide users through structured interactions reducing the complexity to create them through the use of the LLMs to assist us control the conditional logic.

With the mix of prompts, memory management, and power calling we are able to create intuitive and effective dialogue systems, opening recent possibilities for user interaction and task automation.

I hope that this tutorial assist you higher understand tips on how to use LangGraph (I’ve spend a few days banging my head on the wall to grasp how all of the pieces of the library work together).

All of the code of this tutorial will be found here: dmesquita/task_oriented_dialogue_system_langgraph (github.com)

Thanks for reading!

[1] Qin, Libo, et al. “End-to-end task-oriented dialogue: A survey of tasks, methods, and future directions.” arXiv preprint arXiv:2311.09008 (2023).

[2] Prompt generation from user requirements. Available at: https://langchain-ai.github.io/langgraph/tutorials/chatbots/information-gather-prompting