State Transitions Not working as Expected in Agent Setup #167

Closed
opened 2026-02-15 17:16:40 -05:00 by yindo · 0 comments
Owner

Originally created by @jeffersonRosman on GitHub (Feb 6, 2025).

State Transitions Not Displayed Correctly in Mermaid Diagram (see the diagram at the end)

Description:

When generating a mermaid diagram using agent.get_graph().draw_mermaid_png() in LangGraph, the state transitions are not displayed correctly. The transitions do not reflect the defined conditional logic between nodes.

Steps to Reproduce:

Define an AgentState and nodes with transition logic as shown in the code snippet below.
Add nodes and conditional edges to the agent.
Generate the mermaid diagram with agent.get_graph().draw_mermaid_png().
Observe the incorrect relationships in the generated diagram.

Diagram

Image

Code

import logging
from langgraph.graph import MessagesState, StateGraph
from langgraph.checkpoint.memory import MemorySaver
from langgraph.llm import HumanMessage

# Define the AgentState class
class AgentState(MessagesState):
    error: bool
    error_msg: str
    code: str
    external_libs: bool

# Define the nodes
def code_generator(state: AgentState):
    prompt = "Can you provide me with Python code for a function that takes a string as input and returns True if the string is a palindrome and False otherwise? The function should not use any libraries."
    msg = [HumanMessage(content=prompt)]
    function_str = "def is_palindrome(s): return s == s[::-1]"
    return {"messages": [function_str], "code": function_str}

def validate_code(state: AgentState):
    function = state["code"]
    try:
        exec(function)
        return {"error": False}
    except Exception as e:
        return {"error": True, "error_msg": str(e)}

def validate_external_libs(state: AgentState):
    return {"external_libs": False}

def install_external_libs(state: AgentState):
    return {"external_libs": False}

def correct_python_erros(state: AgentState):
    corrected_function_str = "def is_palindrome(s): return s == s[::-1]"
    return {"messages": [corrected_function_str], "code": corrected_function_str, "error": False}

def execute_python_function(state: AgentState):
    function = state["code"]
    exec(function)
    return "Done"

# Define the edges
def check_code(state: AgentState):
    error = state.get("error", None)
    if error:
        return "correct_python_erros"
    else:
        return "validate_external_libs"

def check_external_libs(state: AgentState):
    external_libs = state.get("external_libs", False)
    if external_libs:
        return "install_external_libs"
    else:
        return "execute_python_function"

# Initialize the agent and pass the schema
agent_init = StateGraph(AgentState)
agent_init.add_node("code_generator", code_generator)
agent_init.add_node("validate_code", validate_code)
agent_init.add_node("validate_external_libs", validate_external_libs)
agent_init.add_node("install_external_libs", install_external_libs)
agent_init.add_node("correct_python_erros", correct_python_erros)
agent_init.add_node("execute_python_function", execute_python_function)

# Add the edges
from langgraph.graph import START, END
agent_init.add_edge(START, "code_generator")
agent_init.add_edge("code_generator", "validate_code")
agent_init.add_conditional_edges("validate_code", check_code)
agent_init.add_edge("correct_python_erros", "validate_code")
agent_init.add_conditional_edges("validate_external_libs", check_external_libs)
agent_init.add_edge("install_external_libs", "execute_python_function")
agent_init.add_edge("execute_python_function", END)

# Create the agent
memory = MemorySaver()
agent = agent_init.compile(checkpointer=memory)

# Generate and display the diagram
from IPython.display import Image, display
display(Image(agent.get_graph().draw_mermaid_png()))


Originally created by @jeffersonRosman on GitHub (Feb 6, 2025). State Transitions Not Displayed Correctly in Mermaid Diagram (see the diagram at the end) # Description: When generating a mermaid diagram using agent.get_graph().draw_mermaid_png() in LangGraph, the state transitions are not displayed correctly. The transitions do not reflect the defined conditional logic between nodes. # Steps to Reproduce: Define an AgentState and nodes with transition logic as shown in the code snippet below. Add nodes and conditional edges to the agent. Generate the mermaid diagram with agent.get_graph().draw_mermaid_png(). Observe the incorrect relationships in the generated diagram. # Diagram <img width="849" alt="Image" src="https://github.com/user-attachments/assets/898d21c9-4ece-4eb5-9b62-570ff37c5d4b" /> # Code ```python import logging from langgraph.graph import MessagesState, StateGraph from langgraph.checkpoint.memory import MemorySaver from langgraph.llm import HumanMessage # Define the AgentState class class AgentState(MessagesState): error: bool error_msg: str code: str external_libs: bool # Define the nodes def code_generator(state: AgentState): prompt = "Can you provide me with Python code for a function that takes a string as input and returns True if the string is a palindrome and False otherwise? The function should not use any libraries." msg = [HumanMessage(content=prompt)] function_str = "def is_palindrome(s): return s == s[::-1]" return {"messages": [function_str], "code": function_str} def validate_code(state: AgentState): function = state["code"] try: exec(function) return {"error": False} except Exception as e: return {"error": True, "error_msg": str(e)} def validate_external_libs(state: AgentState): return {"external_libs": False} def install_external_libs(state: AgentState): return {"external_libs": False} def correct_python_erros(state: AgentState): corrected_function_str = "def is_palindrome(s): return s == s[::-1]" return {"messages": [corrected_function_str], "code": corrected_function_str, "error": False} def execute_python_function(state: AgentState): function = state["code"] exec(function) return "Done" # Define the edges def check_code(state: AgentState): error = state.get("error", None) if error: return "correct_python_erros" else: return "validate_external_libs" def check_external_libs(state: AgentState): external_libs = state.get("external_libs", False) if external_libs: return "install_external_libs" else: return "execute_python_function" # Initialize the agent and pass the schema agent_init = StateGraph(AgentState) agent_init.add_node("code_generator", code_generator) agent_init.add_node("validate_code", validate_code) agent_init.add_node("validate_external_libs", validate_external_libs) agent_init.add_node("install_external_libs", install_external_libs) agent_init.add_node("correct_python_erros", correct_python_erros) agent_init.add_node("execute_python_function", execute_python_function) # Add the edges from langgraph.graph import START, END agent_init.add_edge(START, "code_generator") agent_init.add_edge("code_generator", "validate_code") agent_init.add_conditional_edges("validate_code", check_code) agent_init.add_edge("correct_python_erros", "validate_code") agent_init.add_conditional_edges("validate_external_libs", check_external_libs) agent_init.add_edge("install_external_libs", "execute_python_function") agent_init.add_edge("execute_python_function", END) # Create the agent memory = MemorySaver() agent = agent_init.compile(checkpointer=memory) # Generate and display the diagram from IPython.display import Image, display display(Image(agent.get_graph().draw_mermaid_png()))
yindo closed this issue 2026-02-15 17:16:40 -05:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: langchain-ai/langgraphjs#167