LANGCHAIN — What Are Reflection Agents?

Understanding Reflection Agents

Reflection agents are an important technique used to improve the quality and success rate of AI systems, particularly Language Model (LLM) agents. By prompting an LLM to reflect on and critique its past actions, reflection agents can lead to better performance and more thoughtful responses. In this article, we will explore three reflection techniques using LangGraph and provide example implementations of each technique using Python code snippets.

1. Basic Reflection

The basic reflection technique involves composing two LLM calls: a generator and a reflector. The generator responds directly to the user’s requests, while the reflector role plays as a teacher and offers constructive criticism for the initial response. The following Python code snippet demonstrates how this reflection loop can be defined using LangGraph:

from langgraph.graph import MessageGraph

builder = MessageGraph()
builder.add_node("generate", generation_node)
builder.add_node("reflect", reflection_node)
builder.set_entry_point("generate")

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

builder.add_conditional_edges("generate", should_continue)
builder.add_edge("reflect", "generate")
graph = builder.compile()

In this reflection loop, the MessageGraph represents a stateful graph, and the loop proceeds a fixed number of times before returning the final generated output from the generator node.

2. Reflexion

Reflexion, as described by Shinn et. al., is an architecture designed to learn through verbal feedback and self-reflection. This technique involves the actor agent explicitly critiquing each response and grounding its criticism in external data. The following Python code snippet outlines how the reflexion actor loop can be defined using LangGraph:

from langgraph.graph import END, MessageGraph

MAX_ITERATIONS = 5
builder = MessageGraph()
builder.add_node("draft", first_responder.respond)
builder.add_node("execute_tools", execute_tools)
builder.add_node("revise", revisor.respond)
builder.add_edge("draft", "execute_tools")
builder.add_edge("execute_tools", "revise")

def event_loop(state: List[BaseMessage]) -> str:
    num_iterations = _get_num_iterations(state)
    if num_iterations > MAX_ITERATIONS:
        return END
    return "execute_tools"

builder.add_conditional_edges("revise", event_loop)
builder.set_entry_point("draft")
graph = builder.compile()

In this reflexion actor loop, the responder is tasked with generating a response, along with additional actions in the form of search queries, and the revisor is prompted to reflect on the current state.

3. Language Agent Tree Search

Language Agent Tree Search (LATS), introduced by Zhou et. al., is a general LLM agent search algorithm that combines reflection/evaluation and search to achieve better overall task performance. The search process involves selecting the best next actions, expanding and simulating potential actions, reflecting and evaluating outcomes, and backpropagating the scores. The following Python code snippet demonstrates the implementation of LATS using LangGraph:

from langgraph.graph import END, StateGraph

class Node:
    # Node definition

class TreeState(TypedDict):
    # Tree state definition

def should_loop(state: TreeState):
    # Determine whether to continue the tree search

builder = StateGraph(TreeState)
# Graph definition
graph = builder.compile()

In this LATS graph, the selection step picks the node with the highest upper confidence bound (UCT) to balance expected reward and exploration of new paths.

Conclusion

In conclusion, reflection agents play a crucial role in improving the performance and quality of AI systems, particularly LLM agents. The examples provided in this article leverage additional LLM inference to increase the likelihood of generating higher quality output. These techniques can be valuable in scenarios where output quality matters more than response time, and the trajectories can be saved for future model improvement.

By using the provided Python code snippets and examples, you can explore and implement these reflection techniques in your own AI systems to achieve better performance and more thoughtful responses. For detailed implementation and further exploration, you can refer to the LangGraph repository.

By providing clear and detailed Python code snippets and examples, this tutorial-style article aims to help readers understand and implement reflection agents in their own AI systems using the LangGraph framework. The provided code snippets and explanations demonstrate the practical application of reflection techniques, allowing readers to enhance the performance and quality of their AI systems.