Build an AI Coding Agent with LangGraph by LangChain

A D VA NC E D G E NE RAT I VE A I G I T HUB LA NG C HA I N LA RG E LA NG UA G E M O D E LS

Introduction

There has been a massive surge in applications using AI coding agents. With the increasing quality of
LLMs and decreasing cost of inference, it’s only getting easier to build capable AI agents. On top of this,
the tooling ecosystem is evolving rapidly, making it easier to build complex AI coding agents. The
Langchain framework has been a leader on this front. It has all the necessary tools and techniques to
create production-ready AI applications.

But so far, it was lacking in one thing. And that is a multi-agent collaboration with cyclicity. This is crucial
for solving complex problems, where the problem can be divided and delegated to specialized agents. This
is where LangGraph comes into the picture, a part of the Langchain framework designed to accommodate
multi-actor stateful collaboration among AI coding agents. Further, in this article, we will discuss
LangGraph and its basic building blocks while we build an agent with it.

Learning Objectives

Understand what LangGraph is.

Explore the basics of LangGraph for building stateful Agents.
Explore TogetherAI to access open-access models like DeepSeekCoder.
Build an AI coding agent using LangGraph to write unit tests.
This ar ticle was published as a par t of the Data Science Blogathon.

Table of contents
 Introduction
What is LangGraph?
Install Dependencies
Set up LLM

Define Agent State

Define Nodes
Execute the Workflow

LangChain vs LangGraph
Why use LangGraph?
Real-life Use Cases

Conclusion
Frequently Asked Questions

What is LangGraph?

LangGraph is an extension of the LangChain ecosystem. While LangChain allows building AI coding agents
that can use multiple tools to execute tasks, it cannot coordinate multiple chains or actors across the
steps. This is crucial behavior for creating agents that accomplish complex tasks. LangGraph was
conceived keeping these things in mind. It treats the Agent workflows as a cyclic Graph structure, where
each node represents a function or a Langchain Runnable object, and edges are connections between
nodes.

LangGraph’s main features include

Nodes: Any function or Langchain Runnable object like a tool.

Edges: Defines the direction between nodes.
Stateful Graphs: The primary type of graph. It is designed to manage and update state objects as it
processes data through its nodes.

LangGraph leverages this to facilitate a cyclic LLM call execution with state persistence, which is crucial
for agentic behavior. The architecture derives inspiration from Pregel and Apache Beam.

In this article, we will build an Agent for writing Pytest unit tests for a Python class with methods. And this
is the workflow.
We will discuss the concepts in detail as we build our AI coding agent for writing simple unit tests. So, let’s
get to the coding part.

But before that, let’s set up our development environment.

Install Dependencies

First thing first. As with any Python project, create a virtual environment and activate it.

python -m venv auto-unit-tests-writer cd auto-unit-tests-writer source bin/activate

Now, install the dependencies.

!pip install langgraph langchain langchain_openai colorama

Import all the libraries and their classes.

from typing import TypedDict, List import colorama import os from langchain_openai import ChatOpenAI from

langchain_core.messages import SystemMessage from langchain_core.messages import HumanMessage from

langchain_core.runnables import RunnableConfig from langgraph.graph import StateGraph, END from

langgraph.pregel import GraphRecursionError

We will also want to create the directories and files for test cases. You can manually create files or use
Python for that.

# Define the paths. search_path = os.path.join(os.getcwd(), "app") code_file = os.path.join(search_path,

"src/crud.py") test_file = os.path.join(search_path, "test/test_crud.py") # Create the folders and files if

necessary. if not os.path.exists(search_path): os.mkdir(search_path) os.mkdir(os.path.join(search_path,

"src")) os.mkdir(os.path.join(search_path, "test"))

Now, update the crud.py file with code for an in-memory CRUD app. We will use this piece of code to write
unit tests. You can use your Python program for this. We will add the program below to our code.py file.

#crud.py code = """class Item: def __init__(self, id, name, description=None): self.id = id self.name = name
self.description = description def __repr__(self): return f"Item(id={self.id}, name={self.name}, description=

{self.description})" class CRUDApp: def __init__(self): self.items = [] def create_item(self, id, name,

description=None): item = Item(id, name, description) self.items.append(item) return item def read_item(self,
id): for item in self.items: if item.id == id: return item return None def update_item(self, id, name=None,

description=None): for item in self.items: if item.id == id: if name: item.name = name if description:
item.description = description return item return None def delete_item(self, id): for index, item in

enumerate(self.items): if item.id == id: return self.items.pop(index) return None def list_items(self):

return self.items""" with open(code_file, 'w') as f: f.write(code)

Set up LLM

Now, we will specify the LLM we will use in this project. Which model to use here depends on the tasks and
availability of resources. You can use proprietary, powerful models like GPT-4, Gemini Ultra, or GPT-3.5.
Also, you can use open-access models like Mixtral and Llama-2. In this case, as it involves writing codes,
we can use a fine-tuned coding model like DeepSeekCoder-33B or Llama-2 coder. Now, there are multiple
platforms for LLM inferencing, like Anayscale, Abacus, and Together. We will use Together AI to infer
DeepSeekCoder. So, get an API key from Together before going ahead.

from langchain_openai import ChatOpenAI llm = ChatOpenAI(base_url="https://api.together.xyz/v1",

api_key="your-key", model="deepseek-ai/deepseek-coder-33b-instruct")

As together API is compatible with OpenAI SDK, we can use Langchain’s OpenAI SDK to communicate with
models hosted on Together by changing the base_url parameter to “https://api.together.xyz/v1”. In api_key,
pass your Together API key, and in place of models, pass the model name available on Together.

Take your AI innovations to the next level with GenAI Pinnacle. Fine-tune models like Gemini and unlock
endless possibilities in NLP, image generation, and more. Dive in today! Explore Now

Define Agent State

This is one of the crucial parts of LangGraph. Here, we will define an AgentState, responsible for keeping
track of the states of Agents throughout the execution. This is primarily a TypedDict class with entities
that maintain the state of the Agents. Let’s define our AgentState

class AgentState(TypedDict): class_source: str class_methods: List[str] tests_source: str

In the above AgentState class, the class_source stores the original Python class, class_methods for storing
methods of the class, and tests_source for unit test codes. We defined these as AgentState to use them
across execution steps.
Now, define the Graph with the AgentState.

# Create the graph. workflow = StateGraph(AgentState)

As mentioned earlier, this is a stateful graph, and now we have added our state object.

Define Nodes

Now that we have defined the AgentState, we need to add nodes. So, what exactly are nodes? In
LangGraph, nodes are functions or any runnable object, like Langchain tools, that perform a single action.
In our case, we can define several nodes, like a function for finding class methods, a function for inferring
and updating unit tests to state objects, and a function for writing it to a test file.

We also need a way to extract codes from an LLM message. Here’s how.

def extract_code_from_message(message): lines = message.split("\n") code = "" in_code = False for line in
lines: if "```" in line: in_code = not in_code elif in_code: code += line + "\n" return code

The code snippet here assumes the codes to be inside the triple quotes.

Now, let’s define our nodes.

import_prompt_template = """Here is a path of a file with code: {code_file}. Here is the path of a file with
tests: {test_file}. Write a proper import statement for the class in the file. """ # Discover the class and
its methods. def discover_function(state: AgentState): assert os.path.exists(code_file) with open(code_file,

"r") as f: source = f.read() state["class_source"] = source # Get the methods. methods = [] for line in
source.split("\n"): if "def " in line: methods.append(line.split("def ")[1].split("(")[0])
state["class_methods"] = methods # Generate the import statement and start the code. import_prompt =
import_prompt_template.format( code_file=code_file, test_file=test_file ) message =

llm.invoke([HumanMessage(content=import_prompt)]).content code = extract_code_from_message(message)

state["tests_source"] = code + "\n\n" return state # Add a node to for discovery. workflow.add_node(
"discover", discover_function )

In the above code snippet, we defined a function for discovering codes. It extracts the codes from the
AgentState class_source element, dissects the class into individual methods, and passes it to the LLM with
prompts. The output is stored in the AgentState’s tests_source element. We only make it write import
statements for the unit test cases.

We also added the first node to the StateGraph object.

Now, onto the next node.

Also, we can set up some prompt templates that we will need here. These are sample templates you can
change as per your needs.

# System message template. system_message_template = """You are a smart developer. You can do this! You will
write unit tests that have a high quality. Use pytest. Reply with the source code for the test only. Do not

include the class in your response. I will add the imports myself. If there is no test to write, reply with
"# No test to write" and nothing more. Do not include the class in your response. Example: ``` def
test_function(): ... ``` I will give you 200 EUR if you adhere to the instructions and write a high quality

test. Do not write test classes, only methods. """ # Write the tests template. write_test_template = """Here
is a class: ''' {class_source} ''' Implement a test for the method \"{class_method}\". """

Now, define the node.

# This method will write a test. def write_tests_function(state: AgentState): # Get the next method to write

a test for. class_method = state["class_methods"].pop(0) print(f"Writing test for {class_method}.") # Get the
source code. class_source = state["class_source"] # Create the prompt. write_test_prompt =
write_test_template.format( class_source=class_source, class_method=class_method ) print(colorama.Fore.CYAN +
write_test_prompt + colorama.Style.RESET_ALL) # Get the test source code. system_message =

SystemMessage(system_message_template) human_message = HumanMessage(write_test_prompt) test_source =

llm.invoke([system_message, human_message]).content test_source = extract_code_from_message(test_source)
print(colorama.Fore.GREEN + test_source + colorama.Style.RESET_ALL) state["tests_source"] += test_source +

"\n\n" return state # Add the node. workflow.add_node( "write_tests", write_tests_function )

Here, we will make the LLM write test cases for each method, update them to the AgentState’s tests_source
element, and add them to the workflow StateGraph object.

Edges

Now that we have two nodes, we will define edges between them to specify the direction of execution
between them. The LangGraph provides primarily two types of edges.

Conditional Edge: The flow of execution depends on the agents’ response. This is crucial for adding
cyclicity to the workflows. The agent can decide which nodes to move next based on some conditions.
Whether to return to a previous node, repeat the current, or move to the next node.
Normal Edge: This is the normal case, where a node is always called after the invocation of previous
ones.

We do not need a condition to connect discover and write_tests, so we will use a normal edge. Also, define
an entry point that specifies where the execution should start.

# Define the entry point. This is where the flow will start. workflow.set_entry_point("discover") # Always go
from discover to write_tests. workflow.add_edge("discover", "write_tests")

The execution starts with discovering the methods and goes to the function of writing tests. We need
another node to write the unit test codes to the test file.

# Write the file. def write_file(state: AgentState): with open(test_file, "w") as f:

f.write(state["tests_source"]) return state # Add a node to write the file. workflow.add_node( "write_file",
write_file)

As this is our last node, we will define an edge between write_tests and write_file. This is how we can do
this.

# Find out if we are done. def should_continue(state: AgentState): if len(state["class_methods"]) == 0:

return "end" else: return "continue" # Add the conditional edge. workflow.add_conditional_edges(
"write_tests", should_continue, { "continue": "write_tests", "end": "write_file" } )

The add_conditional_edge function takes the write_tests function, a should_continue function that decides
which step to take based on class_methods entries, and a mapping with strings as keys and other
functions as values.

The edge starts at write_tests and, based on the output of should_continue, executes either of the options
in the mapping. For example, if the state[“class_methods”] is not empty, we have not written tests for all
the methods; we repeat the write_tests function, and when we are done writing the tests, the write_file is
executed.
When the tests for all the methods have been inferred from LLM, the tests are written to the test file.

Now, add the final edge to the workflow object for the closure.

# Always go from write_file to end. workflow.add_edge("write_file", END)

Execute the Workflow

The last thing that remained was to compile the workflow and run it.

# Create the app and run it app = workflow.compile() inputs = {} config = RunnableConfig(recursion_limit=100)
try: result = app.invoke(inputs, config) print(result) except GraphRecursionError: print("Graph recursion
limit reached.")

This will invoke the app. The recursion limit is the number of times the LLM will be inferred for a given
workflow. The workflow stops when the limit is exceeded.

You can see the logs on the terminal or in the notebook. This is the execution log for a simple CRUD app.

A lot of the heavy lifting will be done by the underlying model, this was a demo application with the
Deepseek coder model, for better performance you can use GPT-4 or Claude Opus, haiku, etc.
You can also use Langchain tools for web surfing, stock price analysis, etc.

GitHub Gist for complete code: GitHub Gist

Colab Notebook: LangGraph Code Agent

LangChain vs LangGraph

Now, the question is when to use LangChain vs LangGraph.

If the goal is to create a multi-agent system with coordination among them, LangGraph is the way to go.
However, if you want to create DAGs or chains to complete tasks, the LangChain Expression Language is
best suited.

Why use LangGraph?

LangGraph is a potent framework that can improve many existing solutions.

Improve RAG pipelines: LangGraph can augment the RAG with its cyclic graph structure. We can
introduce a feedback loop to evaluate the quality of the retrieved object and, if needed, can improve the
query and repeat the process.
Multi-Agent Workflows: LangGraph is designed to support multi-agent workflows. This is crucial for
solving complex tasks divided into smaller sub-tasks. Different agents with a shared state and
different LLMs and tools can collaborate to solve a single task.
Human-in-the-loop: LangGraph has built-in support for Human-in-the-loop workflow. This means a
human can review the states before moving to the next node.
Planning Agent: LangGraph is well suited to build planning agents, where an LLM planner plans and
decomposes a user request, an executor invokes tools and functions, and the LLM synthesizes
answers based on previous outputs.
Multi-modal Agents: LangGraph can build multi-modal agents, like vision-enabled web navigators.

Real-life Use Cases

There are numerous fields where complex AI coding agents can be helpful.

1. Personal Agents: Imagine having your own Jarvis-like assistant on your electronic devices, ready to
help with tasks at your command, whether it is through text, voice, or even a gesture. That’s one of the
most exciting uses of AI agents!
2. AI Instructors: Chatbots are great, but they have their limits. AI agents equipped with the right tools
can go beyond basic conversations. Virtual AI instructors who can adapt their teaching methods based
on user feedback can be game-changing.
3. Software UX: The user experience of software can be improved with AI agents. Instead of manually
navigating applications, agents can accomplish tasks with voice or gesture commands.
4. Spatial Computing: As AR/VR technology grows in popularity, the demand for AI agents will grow. The
agents can process surrounding information and execute tasks on demand. This may be one of the
best use cases of AI agents shortly.
5. LLM OS: AI-first operating systems where agents are first-class citizens. Agents will be responsible for
doing mundane to complex tasks.

Conclusion
LangGraph is an efficient framework for building cyclic stateful multi-actor agent systems. It fills in the gap
in the original LangChain framework. As it is an extension of LangChain, we can benefit from all the good
things of the LangChain ecosystem. As the quality and capability of LLMs grow, it will be much easier to
create agent systems for automating complex workflows. So, here are the key takeaways from the article.

Key Takeaways

LangGraph is an extension of LangChain, which allows us to build cyclic, stateful, multi-actor agent
systems.
It implements a graph structure with nodes and edges. The nodes are functions or tools, and the edges
are the connections between nodes.
Edges are of two types: conditional and normal. Conditional edges have conditions while going from
one to another, which is important for adding cyclicity to the workflow.
LangGraph is preferred for building cyclic multi-actor agents, while LangChain is better at creating
chains or directed acyclic systems.

Frequently Asked Questions

Q1. What is LangGraph?

Ans. LangGraph is an open-source library for building stateful cyclic multi-actor agent systems. It is built
on top of the LangChain eco-system.

Q2. When to use LangGraph over LangChain?

Ans. LangGraph is preferred for building cyclic multi-actor agents, while LangChain is better at creating
chains or directed acyclic systems.

Q3. What is an AI agent?

Ans. AI agents are software programs that interact with their environment, make decisions, and act to
achieve an end goal.

Q4. What is the best LLM to use with AI agents?

Ans. This depends on your use cases and budget. GPT 4 is the most capable but expensive. For coding,
DeepSeekCoder-33b is a great cheaper option.

Q5. What is the difference between chains and agents?

Ans. The chains are a sequence of hard-coded actions to follow, while agents use LLMs and other tools
(also chains) to reason and act according to the information

The media shown in this ar ticle is not owned by Analytics Vidhya and is used at the Author’s discretion.

Dive into the future of AI with GenAI Pinnacle. From training bespoke models to tackling real-world
challenges like PII masking, empower your projects with cutting-edge capabilities. Star t Exploring.

Article Url - https://www.analyticsvidhya.com/blog/2024/03/build-an-ai-coding-agent-with-langgraph-by-

langchain/

Sunil Kumar Dash