Crafting Smarter AI: Building Self-Improving Reflection Agents with LangGraph

Ahmed Mohiuddin
6 min readOct 21, 2024

--

Introduction

LangGraph is a powerful framework designed for creating stateful, multi-actor applications using large language models (LLMs). It stands out by providing key benefits essential for agent-based workflows: cycles, controllability, and persistence. These features make LangGraph a prime choice for building robust reflection agents, especially in comparison to other frameworks that rely on DAG-based (Directed Acyclic Graph) structures.

Why Use LangGraph for Agent Workflows?

  1. Cycles: Many agent architectures require iterative processes — whether it’s repeating steps or re-evaluating decisions. LangGraph supports these cycles naturally, something that DAG-based solutions can’t handle efficiently.
  2. Controllability: Being a low-level framework, LangGraph gives you precise control over both the flow and the state of your application. This level of customization is crucial when building agents that need to adapt and self-manage.
  3. Persistence: With built-in persistence features, LangGraph enables advanced functionalities like human-in-the-loop interactions and memory, making your agents smarter and more reliable over time.

Reflection

Reflection is a prompting technique designed to enhance the performance and reliability of AI agents. It involves asking a language model to analyze and critique its previous actions, enabling it to identify areas for improvement. This process can also include external data, such as insights from tool interactions, to provide a more informed and thorough reflection, ultimately leading to more effective and refined outputs.

In this blog post, we’ll explore how we can build a simple reflection agent using LangGraph, demonstrating how they can be applied across different domains — one of which is content generation refinement. We will walk through how a reflection agent can critique and improve blog posts in a structured and iterative manner.

Architecture

Simple Reflection Agent
Simple Reflection Agent

As illustrated in the architecture above, when a query is sent to the LLM, it first generates an initial response. This response is then fed back into the LLM for reflection, where it evaluates and critiques its own output. Finally, the LLM uses this critique to produce a refined version of the original response, incorporating the suggested improvements for a more polished result.

Implementation

Now, let’s dive into the code that implements the basic reflection agent to gain a clearer understanding of how it works in practice.

In the .env file, we configure key environment variables that allow our reflection agent to interact with Langchain’s API:

  • LANGCHAIN_TRACING_V2=true: Enables tracing, so you can track and monitor the agent’s behavior using LangSmith.
  • LANGCHAIN_ENDPOINT: Specifies the API endpoint where the LLM’s request and responses can be traced in LangSmith.
  • LANGCHAIN_API_KEY: Your unique API key for LangSmith. This can be created by going to https://smith.langchain.com/
  • LANGCHAIN_PROJECT: Defines the project name ("reflection-agent" in this case) for organizing traces in LangSmith
LANGCHAIN_TRACING_V2=true
LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
LANGCHAIN_API_KEY="YOUR_API_KEY"
LANGCHAIN_PROJECT="reflection-agent"

Here’s the code for the `agent.py` file. It includes two ChatPromptTemplate definitions: one for generating blog content and the other for reflecting and providing feedback.

The MessagesPlaceholder is used to hold the history of messages which contains the generate and reflection feedback messages.

In this example, I’m using the **llama3** model, which is running locally, but you can easily swap it with other models supported by LangChain, such as those from OpenAI, Google, AWS, or any open-source models.

We will create two chains one for generate and other for reflection. These will be invoked as part of the generate and reflection nodes respectively.

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_ollama import OllamaLLM

reflection_prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"You are a seasoned blog editor with a sharp eye for detail. Your task is to evaluate and provide constructive feedback on the user's blog post. "
"Offer an insightful critique focused on areas such as content flow, clarity, depth of information, audience engagement, and overall readability. "
"Make sure to highlight areas where the post can be improved, whether by adjusting length, enhancing structure, or refining tone. "
"Provide actionable recommendations that elevate the quality and impact of the post."
),
MessagesPlaceholder(variable_name="messages")
]
)

generate_content_prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"You are an expert blog-writing assistant entrusted with crafting high-quality, compelling blog posts. "
"Draft the best possible version of the blog post based on the user's input. Focus on clarity, depth, and audience engagement. "
"If the user offers feedback on the draft, thoughtfully revise and improve the post, refining the tone, structure, and content to meet their expectations while maintaining its core message."
),
MessagesPlaceholder(variable_name="messages")
]
)

llm = OllamaLLM(model="llama3")
generate_chain = generate_content_prompt | llm
reflection_chain = reflection_prompt | llm

Now let’s look at the graph we will be building using LangGraph. In addition to the start and end nodes it has the generate and reflect nodes. It also has a conditional edge which decides if we should continue to reflect or return the generated response to the user.

Reflection agent graph
Reflection agent graph

In the code below (main.py), we define two key nodes for our graph: the generate node and the reflection node. State management is handled by passing the state (BaseMessage) between nodes, allowing the flow of data through the process. The output of each node is appended to a list of messages.

In the reflection node implementation below, we wrap the result in a HumanMessage wrapper. This is a clever prompt engineering technique that makes the LLM interpret the reflection or feedback as if it's coming from a human participant in the conversation. By doing this, the model treats the critique more naturally, enhancing its ability to generate meaningful and constructive responses.

Next, we define our graph as a MessageGraph object. We add the necessary nodes, designating the generate node as the entry point. We also establish edges between the nodes, including a conditional edge. In this case, the conditional edge leads to the end node if the state contains more than six messages, guiding the flow of the process based on message count.

The graph is then compiled. The graph can now be invoked by passing in a expert from one of the blog post titled: “Leveraging Emerging Technologies: How Businesses Can Thrive in the Digital Age”

from typing import List,Sequence
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import END, MessageGraph
from agent import generate_chain, reflection_chain

from dotenv import load_dotenv
load_dotenv()

REFLECT = "reflect"
GENERATE = "generate"

def generate_node(state: Sequence[BaseMessage]):
return generate_chain.invoke({"messages": state})


def reflection_node(messages: Sequence[BaseMessage]) -> List[BaseMessage]:
result = reflection_chain.invoke({"messages": messages})
return [HumanMessage(content=result)]


graph = MessageGraph()
graph.add_node(GENERATE, generate_node)
graph.add_node(REFLECT, reflection_node)
graph.set_entry_point(GENERATE)

def should_continue(state: List[BaseMessage]):
if len(state) > 6:
return END
return REFLECT

graph.add_conditional_edges(GENERATE, should_continue)
graph.add_edge(REFLECT, GENERATE)

graph = graph.compile()

if __name__ == "__main__":
inputs = HumanMessage(content="""In today’s fast-paced digital world, businesses must adapt quickly to changing technology trends
in order to stay competitive. Cloud computing, AI, and automation are just a few of the tools that are transforming industries.
However, many companies still struggle with implementing these technologies effectively, often due to a lack of proper planning or expertise.
As a result, they fail to realize the full benefits, such as increased efficiency, cost savings, and better customer experiences.
It’s essential for organizations to not only adopt new technologies but also to invest in training and strategy development to maximize their impact.
""")
response = graph.invoke(inputs)

After we execute the above code we can inspect the output and traces in LangSmith. Below we see the feedback critique provided by reflection.

The generate node then revises the content based on the critique provided as the input.

Conclusion

Reflection agents, built using LangGraph, are a game-changer for workflows that require continuous improvement. Whether refining a blog post, optimizing code, or managing tasks, these agents provide valuable feedback loops that enhance the final output. With LangGraph we can create intelligent, self-improving systems that bring value across many different domains.

If you’re looking to enhance your content creation or build adaptable, iterative agents for other purposes, LangGraph’s reflection agents are an excellent solution to explore.

--

--

Ahmed Mohiuddin
Ahmed Mohiuddin

Written by Ahmed Mohiuddin

Software Engineering Lead and AWS Solutions Architect

No responses yet