From d7b9b3b01d8d3bd4bbd27ad316f695507fd0ddae Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 6 Nov 2024 16:34:29 +0100 Subject: [PATCH 1/2] fix(graph): invalid graph representation if RunnablePassthrough is used --- libs/langgraph/tests/test_pregel.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0f0bbe461..ff56d7394 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -13465,3 +13465,22 @@ def test_debug_nested_subgraphs(): assert stream_task["interrupts"] == history_task.interrupts assert stream_task.get("error") == history_task.error assert stream_task.get("state") == history_task.state + + +def test_runnable_passthrough_node_graph() -> None: + class State(TypedDict): + changeme: str + + async def dummy(state): + return state + + agent = dummy | RunnablePassthrough.assign(prediction=RunnableLambda(lambda x: x)) + + graph_builder = StateGraph(State) + + graph_builder.add_node("agent", agent) + graph_builder.add_edge(START, "agent") + + graph = graph_builder.compile() + + assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() From 4dc08cd724a4069627c83c0949a51a16dcd5b2c1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 6 Nov 2024 08:07:48 -0800 Subject: [PATCH 2/2] Fix --- libs/langgraph/langgraph/graph/graph.py | 18 +- .../tests/__snapshots__/test_pregel.ambr | 606 ++---------------- 2 files changed, 72 insertions(+), 552 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index afe6604dd..e91ac4a47 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -573,18 +573,12 @@ class CompiledGraph(Pregel): metadata["__interrupt"] = "before" elif key in self.interrupt_after_nodes: metadata["__interrupt"] = "after" - if xray: - subgraph = ( - subgraphs[key].get_graph( - config=config, - xray=xray - 1 - if isinstance(xray, int) - and not isinstance(xray, bool) - and xray > 0 - else xray, - ) - if key in subgraphs - else node.get_graph(config=config) + if xray and key in subgraphs: + subgraph = subgraphs[key].get_graph( + config=config, + xray=xray - 1 + if isinstance(xray, int) and not isinstance(xray, bool) and xray > 0 + else xray, ) subgraph.trim_first_node() subgraph.trim_last_node() diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index c690081a2..4268a310a 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -840,65 +840,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 3, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "language_models", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 6, + "id": "agent", "type": "runnable", "data": { "id": [ "langchain", "schema", "runnable", - "RunnablePassthrough" + "RunnableAssign" ], - "name": "Passthrough" + "name": "agent" } }, { @@ -926,46 +877,22 @@ } ], "edges": [ - { - "source": 3, - "target": 4 - }, - { - "source": 4, - "target": 5 - }, - { - "source": 1, - "target": 3 - }, - { - "source": 5, - "target": 2 - }, - { - "source": 1, - "target": 6 - }, - { - "source": 6, - "target": 2 - }, { "source": "__start__", - "target": 1 + "target": "agent" }, { "source": "tools", - "target": 1 + "target": "agent" }, { - "source": 2, + "source": "agent", "target": "tools", "data": "continue", "conditional": true }, { - "source": 2, + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -977,16 +904,10 @@ # name: test_conditional_graph[duckdb].4 ''' graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> agent_parser; - Parallel_agent_outcome_Input --> PromptTemplate; - agent_parser --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -.  continue  .-> tools; - Parallel_agent_outcome_Output -.  exit  .-> __end__; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; ''' # --- @@ -1194,65 +1115,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 3, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "language_models", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 6, + "id": "agent", "type": "runnable", "data": { "id": [ "langchain", "schema", "runnable", - "RunnablePassthrough" + "RunnableAssign" ], - "name": "Passthrough" + "name": "agent" } }, { @@ -1280,46 +1152,22 @@ } ], "edges": [ - { - "source": 3, - "target": 4 - }, - { - "source": 4, - "target": 5 - }, - { - "source": 1, - "target": 3 - }, - { - "source": 5, - "target": 2 - }, - { - "source": 1, - "target": 6 - }, - { - "source": 6, - "target": 2 - }, { "source": "__start__", - "target": 1 + "target": "agent" }, { "source": "tools", - "target": 1 + "target": "agent" }, { - "source": 2, + "source": "agent", "target": "tools", "data": "continue", "conditional": true }, { - "source": 2, + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -1331,16 +1179,10 @@ # name: test_conditional_graph[memory].4 ''' graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> agent_parser; - Parallel_agent_outcome_Input --> PromptTemplate; - agent_parser --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -.  continue  .-> tools; - Parallel_agent_outcome_Output -.  exit  .-> __end__; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; ''' # --- @@ -1548,65 +1390,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 3, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "language_models", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 6, + "id": "agent", "type": "runnable", "data": { "id": [ "langchain", "schema", "runnable", - "RunnablePassthrough" + "RunnableAssign" ], - "name": "Passthrough" + "name": "agent" } }, { @@ -1634,46 +1427,22 @@ } ], "edges": [ - { - "source": 3, - "target": 4 - }, - { - "source": 4, - "target": 5 - }, - { - "source": 1, - "target": 3 - }, - { - "source": 5, - "target": 2 - }, - { - "source": 1, - "target": 6 - }, - { - "source": 6, - "target": 2 - }, { "source": "__start__", - "target": 1 + "target": "agent" }, { "source": "tools", - "target": 1 + "target": "agent" }, { - "source": 2, + "source": "agent", "target": "tools", "data": "continue", "conditional": true }, { - "source": 2, + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -1685,16 +1454,10 @@ # name: test_conditional_graph[postgres].4 ''' graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> agent_parser; - Parallel_agent_outcome_Input --> PromptTemplate; - agent_parser --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -.  continue  .-> tools; - Parallel_agent_outcome_Output -.  exit  .-> __end__; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; ''' # --- @@ -1902,65 +1665,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 3, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "language_models", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 6, + "id": "agent", "type": "runnable", "data": { "id": [ "langchain", "schema", "runnable", - "RunnablePassthrough" + "RunnableAssign" ], - "name": "Passthrough" + "name": "agent" } }, { @@ -1988,46 +1702,22 @@ } ], "edges": [ - { - "source": 3, - "target": 4 - }, - { - "source": 4, - "target": 5 - }, - { - "source": 1, - "target": 3 - }, - { - "source": 5, - "target": 2 - }, - { - "source": 1, - "target": 6 - }, - { - "source": 6, - "target": 2 - }, { "source": "__start__", - "target": 1 + "target": "agent" }, { "source": "tools", - "target": 1 + "target": "agent" }, { - "source": 2, + "source": "agent", "target": "tools", "data": "continue", "conditional": true }, { - "source": 2, + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -2039,16 +1729,10 @@ # name: test_conditional_graph[postgres_pipe].4 ''' graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> agent_parser; - Parallel_agent_outcome_Input --> PromptTemplate; - agent_parser --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -.  continue  .-> tools; - Parallel_agent_outcome_Output -.  exit  .-> __end__; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; ''' # --- @@ -2256,65 +1940,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 3, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "language_models", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 6, + "id": "agent", "type": "runnable", "data": { "id": [ "langchain", "schema", "runnable", - "RunnablePassthrough" + "RunnableAssign" ], - "name": "Passthrough" + "name": "agent" } }, { @@ -2342,46 +1977,22 @@ } ], "edges": [ - { - "source": 3, - "target": 4 - }, - { - "source": 4, - "target": 5 - }, - { - "source": 1, - "target": 3 - }, - { - "source": 5, - "target": 2 - }, - { - "source": 1, - "target": 6 - }, - { - "source": 6, - "target": 2 - }, { "source": "__start__", - "target": 1 + "target": "agent" }, { "source": "tools", - "target": 1 + "target": "agent" }, { - "source": 2, + "source": "agent", "target": "tools", "data": "continue", "conditional": true }, { - "source": 2, + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -2393,16 +2004,10 @@ # name: test_conditional_graph[postgres_pool].4 ''' graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> agent_parser; - Parallel_agent_outcome_Input --> PromptTemplate; - agent_parser --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -.  continue  .-> tools; - Parallel_agent_outcome_Output -.  exit  .-> __end__; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; ''' # --- @@ -2610,65 +2215,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 3, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "language_models", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 6, + "id": "agent", "type": "runnable", "data": { "id": [ "langchain", "schema", "runnable", - "RunnablePassthrough" + "RunnableAssign" ], - "name": "Passthrough" + "name": "agent" } }, { @@ -2696,46 +2252,22 @@ } ], "edges": [ - { - "source": 3, - "target": 4 - }, - { - "source": 4, - "target": 5 - }, - { - "source": 1, - "target": 3 - }, - { - "source": 5, - "target": 2 - }, - { - "source": 1, - "target": 6 - }, - { - "source": 6, - "target": 2 - }, { "source": "__start__", - "target": 1 + "target": "agent" }, { "source": "tools", - "target": 1 + "target": "agent" }, { - "source": 2, + "source": "agent", "target": "tools", "data": "continue", "conditional": true }, { - "source": 2, + "source": "agent", "target": "__end__", "data": "exit", "conditional": true @@ -2747,16 +2279,10 @@ # name: test_conditional_graph[sqlite].4 ''' graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> agent_parser; - Parallel_agent_outcome_Input --> PromptTemplate; - agent_parser --> Parallel_agent_outcome_Output; - Parallel_agent_outcome_Input --> Passthrough; - Passthrough --> Parallel_agent_outcome_Output; - __start__ --> Parallel_agent_outcome_Input; - tools --> Parallel_agent_outcome_Input; - Parallel_agent_outcome_Output -.  continue  .-> tools; - Parallel_agent_outcome_Output -.  exit  .-> __end__; + __start__ --> agent; + tools --> agent; + agent -.  continue  .-> tools; + agent -.  exit  .-> __end__; ''' # ---