Building Powerful AI Coding Agents with LangGraph and LangChain

Introduction

The rise of large language models (LLMs) and frameworks like LangChain has sparked massive growth and interest in AI coding agents. These AI assistants can understand natural language instructions, break down complex tasks, leverage external tools, and generate or edit code to help developers be more productive.

As LLMs continue to improve in quality while decreasing in cost, it‘s becoming easier than ever to create AI coding agents and harness their power. However, one key aspect that‘s been missing is seamless multi-agent collaboration.

That‘s where LangGraph comes in. LangGraph is an extension of the popular LangChain framework that introduces stateful graphs to coordinate multiple AI agents in a cyclical workflow. This unlocks the ability to build much more sophisticated coding agents that can tackle complex, multi-step tasks by delegating subtasks to specialized agents.

In this article, we‘ll dive into what LangGraph is, explore its key concepts and building blocks, and walk through an example of using it to create an AI agent that writes unit tests for Python code. By the end, you‘ll have a solid grasp of LangGraph and a blueprint for building your own powerful AI coding assistants.

What is LangGraph?

LangGraph builds upon the foundation of LangChain, a framework for developing applications powered by language models. While LangChain provides reusable components for prompting LLMs, connecting to external tools, storing data, and creating sequential chains of actions, it lacks native support for multi-agent coordination.

LangGraph addresses this by representing agent workflows as cyclic graph structures. Each node in the graph is a function or callable object (like a LangChain tool), while the edges define the flow of data and execution between nodes.

The key features and concepts of LangGraph include:

• Stateful Graphs: LangGraph uses stateful graphs that can persist and update a state object as data flows through the nodes. This allows agents to maintain context and share information.

• Nodes: Nodes are the core building blocks and can be any Python function or LangChain component. They represent individual actions or subtasks in the overall workflow.

• Edges: Edges are the connections between nodes that determine the direction of data flow. They can be normal (always flowing from one node to the next) or conditional (using if/else logic to create cycles and branching paths).

• Cycles and Recursion: By allowing cycles in the graph, LangGraph can create recursive agent workflows. An agent can repeatedly call itself or another agent until a goal is met, like an AI planner refining its outputs.

Under the hood, LangGraph takes inspiration from graph computing models like Pregel and Apache Beam to efficiently process data through the nodes and edges. But from a user‘s perspective, the key idea is that multiple AI functions can be organized in a graph to perform complex tasks.

Building an AI Unit Test Agent with LangGraph

To make things concrete, let‘s walk through an example of using LangGraph to build an AI agent that can automatically write unit tests for a Python codebase. We‘ll break this down step-by-step, explaining the key concepts as we go.

Step 1: Install Dependencies

First, make sure you have Python and pip installed. Then create a new virtual environment and install the required packages:

python -m venv langraph-agent
source langraph-agent/bin/activate
pip install langchain langraph openai tiktoken

This will install LangChain, LangGraph, and some additional dependencies. We‘ll be using OpenAI‘s APIs for the underlying LLM, but you could swap this out for other models.

Step 2: Define the Agent State

Next, we need to define a type that will represent the shared state that our agent functions can access and modify. For our unit test agent, we‘ll use a simple dictionary that stores the source code to test, any existing test code, and the names of functions to test:

from typing import List, Dict

class AgentState(dict):
    source_code: str
    test_code: str
    functions_to_test: List[str]

We inherit from dict so the state can easily be serialized and passed between graph nodes. The source_code field will store the Python code we want to test, test_code will incrementally store the unit tests we generate, and functions_to_test will keep track of which functions still need tests.

Step 3: Initialize the LangGraph

With our agent state defined, we can initialize a LangGraph stateful graph:

from langraph.graph import StateGraph

agent = StateGraph(AgentState)

This creates an empty graph whose nodes will operate on our AgentState type. We‘ll add nodes and edges to this to define the agent workflow.

Step 4: Create Agent Nodes

Now we can define the actual agent functions that make up the nodes in our graph. We‘ll create three nodes:

1. extract_functions: This node will take the initial Python source code, extract the function names that need to be tested, and initialize the agent state.

2. generate_test: This node will take a single function name and the source code, use an LLM to generate a unit test for it, and append that test code to the shared state.

3. save_tests: This node will save the accumulated test code to a file once all functions have been processed.

Here‘s what the implementations look like:

import re

def extract_functions(state: AgentState, source_code: str):
    function_names = re.findall(r‘def (\w+)\(‘, source_code)
    state.update(
        source_code=source_code, 
        test_code=‘import unittest\n‘,
        functions_to_test=function_names
    )
    return state

def generate_test(state: AgentState, openai_api_key: str):
    if not state[‘functions_to_test‘]:
        return state  # No more functions to test

    function_name = state[‘functions_to_test‘].pop(0)

    prompt = f"""
    Here is a function from a Python file:

    {state[‘source_code‘]}

    Write a unit test for the `{function_name}` function using the Python unittest framework.
    Focus on testing expected behavior and edge cases. Don‘t include any other setup code.
    """

    response = openai.Completion.create(
        engine=‘code-davinci-002‘,
        prompt=prompt,
        max_tokens=256,
        temperature=0.7,
        top_p=1,
        n=1,
        stop=None
    )

    test_code = response.choices[0].text.strip()
    state[‘test_code‘] += f‘\n\nclass Test{function_name}(unittest.TestCase):\n{test_code}\n‘

    return state

def save_tests(state: AgentState, output_path: str):
    with open(output_path, ‘w‘) as f:
        f.write(state[‘test_code‘])

    return state

The extract_functions node is pretty straightforward – it takes the source code, finds all function definitions using a regular expression, and initializes the agent state with that code and list of functions.

The generate_test node is where the LLM magic happens. It pops the next function name off the functions_to_test list in the state, constructs a prompt asking the LLM to generate a unit test for that specific function, and appends the result to the accumulated test_code.

Finally, the save_tests node takes the generated test code from the state and simply saves it to a file path.

Step 5: Connect Nodes and Edges

With our agent nodes defined, the last step is to connect them into a workflow using LangGraph edges. The graph we want is pretty simple:

1. Run extract_functions on the initial source code
2. Run generate_test repeatedly until all functions have been processed
3. Run save_tests to write the results to a file

We can define this flow with just a few lines:

agent.add_node(extract_functions)
agent.add_node(generate_test)  
agent.add_node(save_tests)

agent.add_edge(‘extract_functions‘, ‘generate_test‘)
agent.add_conditional_edge(
    ‘generate_test‘, 
    lambda state: len(state[‘functions_to_test‘]) > 0, 
    ‘generate_test‘,
    ‘save_tests‘
)

The key parts here are:

• add_node registers each of our agent functions as graph nodes
• add_edge creates a normal edge between extract_functions and generate_test, ensuring generate_test will run after the initial extraction
• add_conditional_edge is where the magic happens. This creates a conditional edge from generate_test back to itself, so it will run repeatedly as long as there are more functions to process. Once the functions_to_test list in the state is empty, it will proceed to the save_tests node instead.

And with that, our agent is complete! We can run it by invoking the graph on some source code:

source_code = """
def add(x, y):
    return x + y

def subtract(x, y):
    return x - y
"""

result = agent.run(source_code=source_code, openai_api_key=my_api_key, output_path=‘tests.py‘)

print(result[‘test_code‘])

This will generate unit tests for the add and subtract functions, something like:

import unittest

class TestAdd(unittest.TestCase):
    def test_add(self):
        self.assertEqual(add(1, 2), 3)
        self.assertEqual(add(-1, 1), 0)
        self.assertEqual(add(0, 0), 0)

class TestSubtract(unittest.TestCase):  
    def test_subtract(self):
        self.assertEqual(subtract(5, 3), 2)
        self.assertEqual(subtract(-1, 1), -2)
        self.assertEqual(subtract(0, 0), 0)

And that‘s it! We‘ve created an AI coding agent that can read a Python file, automatically generate unit tests using an LLM, and save the results. All powered by LangGraph.

Of course, this is just a toy example – you could extend this with additional nodes for refining the tests, handling more complex types of functions, or many other improvements. But it demonstrates the core concepts of LangGraph and how it can coordinate multiple agent functions in a cyclic workflow.

When to Use LangGraph vs. LangChain

So when should you reach for LangGraph compared to vanilla LangChain? The key differentiators are:

• LangGraph supports stateful cyclic graphs, which makes it a better fit for recursive, multi-step workflows. If your task can be cleanly divided into subtasks that might need to repeat or loop, LangGraph is a great choice.

• LangChain focuses more on linear sequences of actions (chains). If your agent task is straightforward and doesn‘t require looping back to previous steps, LangChain is usually sufficient and simpler.

• LangGraph also puts more emphasis on multi-agent collaboration by making it easy for multiple agent functions to share state. If you foresee your system growing to multiple specialized agents that need to coordinate, LangGraph will make that much easier than bare LangChain.

But in practice, you‘ll likely end up using both in tandem. Most of LangChain‘s components can be used as nodes in a LangGraph, so it‘s easy to combine the two frameworks.

The Future of AI Coding Agents

AI coding agents are one of the most exciting current applications of large language models. The ability to not just converse with an AI assistant but have it actually read, analyze, and write code opens up huge opportunities for developer productivity.

Some future use cases and directions I‘m excited about:

• Smarter code review: AI agents that can deeply analyze a codebase, suggest optimizations and best practices, and even have back-and-forth dialogues with developers about design choices.

• Automated refactoring: Agents that can reliably refactor and modernize large codebases, like upgrading a Python 2 project to Python 3 or migrating from one web framework to another.

• AI pair programming: More interactive development where an AI agent acts as a real-time coding assistant, offering suggestions, catching bugs, and even writing parts of the implementation alongside a human.

• Natural language interfaces: Imagine describing your intended application in plain English, then having an AI agent generate an initial prototype, suggest a test plan, and even deploy it for you. LangChain and LangGraph are great foundations for this kind of end-to-end agent.

The key challenges will be improving the reliability and safety of these agents (nobody wants an AI assistant to introduce subtle bugs!), but I believe those are tractable problems. Tools like LangGraph that let us factor agent workflows into smaller, more understandable pieces are an important part of the solution.

Conclusion and Learning More

To recap, in this article we:

• Introduced the concept of AI coding agents and how they‘re becoming more powerful and accessible thanks to improving LLMs
• Explored LangGraph as an extension of LangChain that supports complex multi-agent workflows using stateful directed graphs
• Walked through a detailed example of building a Python unit test generation agent using LangGraph
• Discussed when to use LangGraph vs. LangChain and the future potential of AI coding assistants

I hope this has given you a practical introduction to LangGraph and a sense of what‘s possible with AI coding agents. The field is evolving incredibly quickly, but the core ideas of chaining together LLM-powered functions, connecting to external tools, and factoring complex tasks into agent workflows will continue to be essential.

If you want to learn more and go deeper with LangGraph, I recommend:

• The official LangGraph documentation, which includes installation instructions, API references, and more detailed guides.
• The LangChain use case guides, which showcase more end-to-end examples of building AI agents (many of which could be extended with LangGraph).
• The LangChain Python docs for details on all the LangChain components you can use with LangGraph.
• Experimenting yourself! The best way to build intuition for these tools is to dive in and try building your own coding agents. Start small, iterate, and see what you can create.

We‘re still in the early stages of the AI agent revolution, but I believe tools like LangGraph are a major step forward in realizing the potential of AI-assisted development. Excited to see what you all build!