Cannot draw graph w/ xray=True when graph contains a subgraph w/ edges targeting the subgraph entry point

[mhmaguire]

Checked other resources

• I added a very descriptive title to this issue.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.
• I am sure that this is a bug in LangChain rather than my code.

Example Code

from typing import TypedDict, Annotated, List
from langchain_core.messages import AnyMessage
from langgraph.graph import StateGraph, add_messages
from IPython.display import Image, display
import random


class State(TypedDict):
    messages: Annotated[List[AnyMessage], add_messages]


def node(name):
    def _node(state: State):
        return {'messages': [('human', f'entered {name} node')]}

    return _node


parent = StateGraph(State)
child = StateGraph(State)

child.add_node('c_one', node('c_one'))
child.add_node('c_two', node('c_two'))

child.add_edge('__start__', 'c_one')
child.add_edge('c_two', 'c_one')

child.add_conditional_edges('c_one',
                            lambda x: str(random.randrange(0, 2)),
                            {
                                '0': 'c_two',
                                '1': '__end__'
                                
                            })

parent.add_node('p_one', node('p_one'))
parent.add_node('p_two', child.compile())

parent.add_edge('__start__', 'p_one')
parent.add_edge('p_two', 'p_one')

parent.add_conditional_edges('p_one',
                            lambda x: str(random.randrange(0, 2)),
                            {
                                '0': 'p_two',
                                '1': '__end__'
                                
                            })


app = parent.compile()

for event in app.stream({'messages': [('human', 'start')]}):
    print(event)


display(Image(app.get_graph(xray=True).draw_mermaid_png())) # this fails

-->

{'p_one': {'messages': [('human', 'entered p_one node')]}}
{'p_two': {'messages': [HumanMessage(content='start', id='53318df0-564d-4f9b-b1ad-a8bcf03e2219'), HumanMessage(content='entered p_one node', id='c67dcfa5-e1cd-4bc0-a9b5-2ca6676150bd'), HumanMessage(content='entered c_one node', id='35b095bb-3113-4d93-8afa-b8225cb72042'), HumanMessage(content='entered c_two node', id='8f231424-9c4f-49d2-a61e-12917d953454'), HumanMessage(content='entered c_one node', id='b8bdc719-ff93-40ef-adf5-c4e90d541a2a')]}}
{'p_one': {'messages': [('human', 'entered p_one node')]}}
{'p_two': {'messages': [HumanMessage(content='start', id='53318df0-564d-4f9b-b1ad-a8bcf03e2219'), HumanMessage(content='entered p_one node', id='c67dcfa5-e1cd-4bc0-a9b5-2ca6676150bd'), HumanMessage(content='entered c_one node', id='35b095bb-3113-4d93-8afa-b8225cb72042'), HumanMessage(content='entered c_two node', id='8f231424-9c4f-49d2-a61e-12917d953454'), HumanMessage(content='entered c_one node', id='b8bdc719-ff93-40ef-adf5-c4e90d541a2a'), HumanMessage(content='entered p_one node', id='8cf2a906-baa2-41c6-84f8-615546a90dcf'), HumanMessage(content='entered c_one node', id='05991dcf-6d6e-4968-bb3a-5f1b39bbac0b')]}}
{'p_one': {'messages': [('human', 'entered p_one node')]}}
{'p_two': {'messages': [HumanMessage(content='start', id='53318df0-564d-4f9b-b1ad-a8bcf03e2219'), HumanMessage(content='entered p_one node', id='c67dcfa5-e1cd-4bc0-a9b5-2ca6676150bd'), HumanMessage(content='entered c_one node', id='35b095bb-3113-4d93-8afa-b8225cb72042'), HumanMessage(content='entered c_two node', id='8f231424-9c4f-49d2-a61e-12917d953454'), HumanMessage(content='entered c_one node', id='b8bdc719-ff93-40ef-adf5-c4e90d541a2a'), HumanMessage(content='entered p_one node', id='8cf2a906-baa2-41c6-84f8-615546a90dcf'), HumanMessage(content='entered c_one node', id='05991dcf-6d6e-4968-bb3a-5f1b39bbac0b'), HumanMessage(content='entered p_one node', id='a67bd7a7-dbb0-418b-913f-dc69f02592a7'), HumanMessage(content='entered c_one node', id='5785cb97-4bcd-4c76-a8d6-bb3e62e967f4')]}}
{'p_one': {'messages': [('human', 'entered p_one node')]}}

Error Message and Stack Trace (if applicable)

see below

Description

I'm trying to draw a graph which contains a subgraph as a node.
The subgraph has edges which target the entry node, which seems to be the problem in get_graph(xray=True)

I expect to see an image of the graph.

Instead I get an exception.

System Info

see below

[mhmaguire]

Ran sys info in the wrong env. Here are actual values for stacktrace and env info:

System Information
------------------
> OS:  Linux
> OS Version:  #35-Ubuntu SMP PREEMPT_DYNAMIC Fri Apr 26 11:23:57 UTC 2024
> Python Version:  3.11.1 (main, Jan  9 2023, 21:29:42) [GCC 12.2.0]

Package Information
-------------------
> langchain_core: 0.2.0
> langchain: 0.2.1
> langchain_community: 0.2.0
> langsmith: 0.1.23
> langchain_anthropic: 0.1.13
> langchain_experimental: 0.0.59
> langchain_openai: 0.1.7
> langchain_postgres: 0.0.6
> langchain_qdrant: 0.1.0
> langchain_text_splitters: 0.2.0
> langchainhub: 0.1.15
> langgraph: 0.0.62

Packages not installed (Not Necessarily a Problem)
--------------------------------------------------
The following packages were not found:

> langserve

{
	"name": "AttributeError",
	"message": "'NoneType' object has no attribute 'id'",
	"stack": "---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 58
     54 for event in app.stream({'messages': [('human', 'start')]}, stream_mode='values'):
     55     print(event)
---> 58 display(Image(app.get_graph(xray=True).draw_mermaid_png())) # this fails

File ~/.pyenv/versions/3.11.1/envs/math/lib/python3.11/site-packages/langgraph/graph/graph.py:470, in CompiledGraph.get_graph(self, config, xray)
    468         end_nodes[key] = n
    469 for start, end in sorted(self.builder._all_edges):
--> 470     graph.add_edge(start_nodes[start], end_nodes[end])
    471 for start, branches in self.builder.branches.items():
    472     default_ends = {
    473         **{k: k for k in self.builder.nodes if k != start},
    474         END: END,
    475     }

File ~/.pyenv/versions/3.11.1/envs/math/lib/python3.11/site-packages/langchain_core/runnables/graph.py:260, in Graph.add_edge(self, source, target, data, conditional)
    252 def add_edge(
    253     self,
    254     source: Node,
   (...)
    257     conditional: bool = False,
    258 ) -> Edge:
    259     \"\"\"Add an edge to the graph and return it.\"\"\"
--> 260     if source.id not in self.nodes:
    261         raise ValueError(f\"Source node {source.id} not in graph\")
    262     if target.id not in self.nodes:

AttributeError: 'NoneType' object has no attribute 'id'"
}

[hinthornw]

Can reproduce - thanks for raising

[jakelucasnyc]

Is there any update on this? I'm still encountering this issue with my graphs

[tecton]

My workaround is to remove the following 2 lines (link):

subgraph.trim_first_node()
subgraph.trim_last_node()

In the above case, the subgraph cannot find first node / last node after they are trimmed.

[RaminZi]

Getting the same error when having a conditional edge in my subgraph while the whole graph runs just fine!

[cigotete]

Same error here, the sub-graph and graph runs ok, but the graphic can not be generated using "xray" parameter. Error message is the same: AttributeError: 'NoneType' object has no attribute 'id'.

[Glinte]

@nfcampos Mind giving this issue a bit of a priority boost? Almost every complex graph have this problem.

[chemyibinjiang]

same problem here

[nfcampos]

Hi, can you confirm this issue still occurs in the latest version? 0.2.4

[nfcampos]

I've added the snippet above as a test case and can confirm this is fixed, closing, thanks for reporting!