DOCSEARCHbyAlgolia
Submit Docs

LangGraph Reflection

repository·main·Indexed 41 days ago

https://github.com/langchain-ai/langgraph-reflection

A package providing prebuilt graph architectures for implementing reflection-style agents. It utilizes a loop between a main solver agent and a critique agent to iteratively improve outputs based on feedback, such as LLM-as-a-judge evaluations or static code analysis via Pyright. Key features include the `create_reflection_graph` utility for combining assistant and judge graphs, and `MessagesWithSteps` for managing iteration limits to prevent infinite loops.

Tokens
2.4K
Snippets
6
Records
7
Agent score
2%

What's inside LangGraph Reflection

  1. How the reflection agent architecture works

    main

    The reflection agent uses a two-agent loop to iteratively improve outputs. It consists of a main agent (the solver) and a critique agent (the evaluator).

    The Workflow:

    1. The main agent is called to solve the user's task.
    2. The critique agent is called to evaluate the main agent's work.
    3. Decision Logic:
      • If the critique agent finds issues, it returns a user message containing the critique. This triggers the main agent to run again to address the feedback.
      • If there is nothing to critique, the reflection agent finishes.
    4. This loop repeats until the critique agent finds no further issues.

    Requirements for your graphs:

    • The main agent must accept a list of messages as input.
    • The critique agent (the judge) must return a user message if critiques exist, otherwise it should return no messages (or None) to signal completion.

  3. How end_or_reflect controls the reflection loop

    main

    The end_or_reflect function is a conditional edge function used within a reflection graph to determine if the process should continue or stop.

    Termination Logic:

    1. Step Limit: If state["remaining_steps"] is less than 2, it returns END.
    2. Empty Messages: If there are no messages in the state, it returns END.
    3. Message Type: It inspects the last message in the state:
      • If the last message is a HumanMessage, it returns "graph" (triggering another round of assistant processing).
      • Otherwise (e.g., if the last message was from the AI/Judge), it returns END.

    This pattern allows a judge to provide feedback that is treated as a new user instruction for the assistant.

    # Logic summary of end_or_reflect:
# if remaining_steps < 2 -> END
# if len(messages) == 0 -> END
# if isinstance(last_message, HumanMessage) -> "graph"
# else -> END

  4. Implement LLM-as-a-Judge reflection

    main

    You can use an LLM to act as the critique agent. The judge evaluates the assistant's response based on criteria like accuracy, completeness, clarity, helpfulness, and safety.

    To implement this, define a judge function that uses an evaluator (e.g., from openevals) and returns a new user message containing the critique if the score is low. Then, use create_reflection_graph to combine your assistant graph with the judge graph.

    # Define the main assistant graph
assistant_graph = ...

# Define the judge function that evaluates responses
def judge_response(state, config):
    """Evaluate the assistant's response using a separate judge model."""
    evaluator = create_llm_as_judge(   
        prompt=critique_prompt,
        model="openai:o3-mini",
        feedback_key="pass",
    )
    eval_result = evaluator(outputs=state["messages"][-1].content, inputs=None)

    if eval_result["score"]:
        print("✅ Response approved by judge")
        return
    else:
        # Otherwise, return the judge's critique as a new user message
        print("⚠️ Judge requested improvements")
        return {"messages": [{"role": "user", "content": eval_result["comment"]}]}

# Create graphs with reflection
judge_graph = StateGraph(MessagesState).add_node(judge_response)....

# Create reflection graph that combines assistant and judge
reflection_app = create_reflection_graph(assistant_graph, judge_graph)
result = reflection_app.invoke({"messages": example_query})

  5. Implement Code Validation reflection using Pyright

    main

    You can use the reflection agent to validate and improve Python code by using Pyright for static type checking.

    In this pattern, the critique agent extracts code from the conversation, runs a Pyright evaluator, and if errors are found, returns a user message containing the error details. This forces the main agent to attempt a fix in the next iteration.

    Required dependencies: pip install langgraph-reflection langchain openevals pyright

    assistant_graph = ...

# Function that validates code using Pyright
def try_running(state: dict) -> dict | None:
    """Attempt to run and analyze the extracted Python code."""
    # Extract code from the conversation
    code = extract_python_code(state['messages'])
    
    # Run Pyright analysis
    evaluator = create_pyright_evaluator()
    result = evaluator(outputs=code)
    
    if not result['score']:
        # If errors found, return critique for the main agent
        return {
            "messages": [{
                "role": "user",
                "content": f"I ran pyright and found this: {result['comment']}\n\n"
                          "Try to fix it..."
            }]
        }
    # No errors found - return None to indicate success
    return None

# Create graphs with reflection
judge_graph = StateGraph(MessagesState).add_node(try_running)....

# Create reflection system that combines code generation and validation
reflection_app = create_reflection_graph(assistant_graph, judge_graph)
result = reflection_app.invoke({"messages": example_query})

  6. Create a reflection graph with create_reflection_graph

    main

    Use create_reflection_graph to wrap an existing assistant graph and a reflection (judge) graph into a single state machine that performs iterative refinement.

    The resulting graph follows this flow:

    1. Starts at the graph node (the assistant).
    2. Moves to the reflection node (the judge).
    3. Uses end_or_reflect to decide whether to loop back to the assistant for refinement or terminate.

    Requirements for state_schema:

    • Must contain a messages key.
    • Must not contain a remaining_steps key (as this is automatically added to prevent shadowing).

    Parameters:

    • graph: The CompiledStateGraph representing the primary assistant.
    • reflection: The CompiledStateGraph representing the judge/evaluator.
    • state_schema (optional): The schema for the state. If not provided, it defaults to the schema of the input graph.
    • config_schema (optional): The configuration schema for the new graph.
    from langgraph_reflection import create_reflection_graph

# Assuming 'assistant_graph' and 'judge_graph' are already compiled
reflection_graph = create_reflection_graph(
    graph=assistant_graph,
    reflection=judge_graph
)

  7. Use MessagesWithSteps for state management in reflection loops

    main

    The MessagesWithSteps class is a specialized version of MessagesState used to track the number of remaining iterations in a reflection loop. It adds a remaining_steps field of type RemainingSteps to the standard message state.

    This is used by the end_or_reflect logic to ensure the agent does not enter an infinite loop of self-correction.

    from langgraph_reflection import MessagesWithSteps

# MessagesWithSteps inherits from MessagesState and adds:
# remaining_steps: RemainingSteps