AI Agents with LangGraph

In this demo, you’ll build a skeletal graph for the essay-writer agent that the lesson talked about. You won’t use any LLMs this time. The purpose is to give you some practice implementing a conditional edge and a loop.

Start by opening the decisions.ipynb notebook, which you’ll find in this lesson’s Starter folder.

Create a writer function:

def writer(input):
  return "an essay"

You could use an LLM to write a full essay based on the input topic, but today, the function will return a fixed string regardless of the input. Then add the reviser function:

revisions = 0

def reviser(input):
  global revisions
  revisions = revisions + 1
  print(f"revision number {revisions}")
  return input + " revisions"

Today, you’ll use a stateless Graph. In a full-fledged application, you’d want to use a StateGraph where the state is passed from node to node. To make up for that, you’ll use a global variable to count the number of revisions the graph makes. Again, you could use an LLM to actually revise the essay, but in this case, you’ll just return some text.

Next, add a checker function:

def checker(input):
  if revisions >= 3:
    return "good"
  else:
    return "feedback"

An LLM-based agent could provide some real feedback points for the reviser to work on. For simplicity, though, you’ll return a fixed string.

Now, import LangGraph and build your nodes.

from langgraph.graph import Graph, START, END

graph = Graph()

graph.add_node("writer", writer)
graph.add_node("reviser", reviser)
graph.add_node("checker", checker)

"writer" is the node name, and writer is the function. "reviser" is the node name, and reviser is the function. "checker" is the node name, and checker is the function.

Then, connect the normal edges between the nodes:

graph.add_edge(START, "writer")
graph.add_edge("writer", "checker")
graph.add_edge("reviser", "checker")

Before you add any conditional edges, you need to define a routing function:

def check(input):
  if input == "good":
    return "pass"
  else:
    return "fail"

The input will come from the checker node. Now you’re ready to add the conditional edges:

graph.add_conditional_edges(
  "checker",
  check,
  {
    "fail": "reviser",
    "pass": END
  }
)

Compile the graph:

app = graph.compile()

There’s a useful function that lets you visualize a compiled graph. Add that next:

from IPython.display import Image, display

display(Image(app.get_graph().draw_mermaid_png()))

The dotted lines show the conditional edges. If the check fails, the graph goes to the reviser. If it passes, the graph goes to END. You can also see the loop going between the checker and the reviser.

The &nbsp around fail and pass shouldn’t be there, but I wasn’t able to get rid of them while making these lessons.

Anyway, that’s how you make a graph with conditional edges and a loop.