core[patch]: pass node name in graph.to_json()

langchain-ai · Nov 1, 2024 · 5dc3db0 · 5dc3db0
1 parent 71f590d
commit 5dc3db0
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 0 deletions.
diff --git a/libs/core/langchain_core/runnables/graph.py b/libs/core/langchain_core/runnables/graph.py
@@ -294,6 +294,7 @@ def to_json(self, *, with_schemas: bool = False) -> dict[str, list[dict[str, Any
             "nodes": [
                 {
                     "id": stable_node_ids[node.id],
+                    "name": node.name,
                     **node_data_json(node, with_schemas=with_schemas),
                 }
                 for node in self.nodes.values()

diff --git a/libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr b/libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr
@@ -177,6 +177,7 @@
       dict({
         'data': 'PromptInput',
         'id': 0,
+        'name': 'PromptInput',
         'type': 'schema',
       }),
       dict({
@@ -190,6 +191,7 @@
           'name': 'PromptTemplate',
         }),
         'id': 1,
+        'name': 'PromptTemplate',
         'type': 'runnable',
       }),
       dict({
@@ -203,16 +205,19 @@
           'name': 'FakeListLLM',
         }),
         'id': 2,
+        'name': 'FakeListLLM',
         'type': 'runnable',
       }),
       dict({
         'data': 'Parallel<as_list,as_str>Input',
         'id': 3,
+        'name': 'Parallel<as_list,as_str>Input',
         'type': 'schema',
       }),
       dict({
         'data': 'Parallel<as_list,as_str>Output',
         'id': 4,
+        'name': 'Parallel<as_list,as_str>Output',
         'type': 'schema',
       }),
       dict({
@@ -226,16 +231,19 @@
           'name': 'CommaSeparatedListOutputParser',
         }),
         'id': 5,
+        'name': 'CommaSeparatedListOutputParser',
         'type': 'runnable',
       }),
       dict({
         'data': 'conditional_str_parser_input',
         'id': 6,
+        'name': 'conditional_str_parser_input',
         'type': 'schema',
       }),
       dict({
         'data': 'conditional_str_parser_output',
         'id': 7,
+        'name': 'conditional_str_parser_output',
         'type': 'schema',
       }),
       dict({
@@ -249,6 +257,7 @@
           'name': 'StrOutputParser',
         }),
         'id': 8,
+        'name': 'StrOutputParser',
         'type': 'runnable',
       }),
       dict({
@@ -262,6 +271,7 @@
           'name': 'XMLOutputParser',
         }),
         'id': 9,
+        'name': 'XMLOutputParser',
         'type': 'runnable',
       }),
     ]),
@@ -331,6 +341,7 @@
           'type': 'object',
         }),
         'id': 0,
+        'name': 'PromptInput',
         'type': 'schema',
       }),
       dict({
@@ -344,6 +355,7 @@
           'name': 'PromptTemplate',
         }),
         'id': 1,
+        'name': 'PromptTemplate',
         'type': 'runnable',
       }),
       dict({
@@ -357,6 +369,7 @@
           'name': 'FakeListLLM',
         }),
         'id': 2,
+        'name': 'FakeListLLM',
         'type': 'runnable',
       }),
       dict({
@@ -1765,6 +1778,7 @@
           'title': 'RunnableParallel<as_list,as_str>Input',
         }),
         'id': 3,
+        'name': 'Parallel<as_list,as_str>Input',
         'type': 'schema',
       }),
       dict({
@@ -1789,6 +1803,7 @@
           'type': 'object',
         }),
         'id': 4,
+        'name': 'Parallel<as_list,as_str>Output',
         'type': 'schema',
       }),
       dict({
@@ -1802,6 +1817,7 @@
           'name': 'CommaSeparatedListOutputParser',
         }),
         'id': 5,
+        'name': 'CommaSeparatedListOutputParser',
         'type': 'runnable',
       }),
       dict({
@@ -1810,13 +1826,15 @@
           'type': 'string',
         }),
         'id': 6,
+        'name': 'conditional_str_parser_input',
         'type': 'schema',
       }),
       dict({
         'data': dict({
           'title': 'conditional_str_parser_output',
         }),
         'id': 7,
+        'name': 'conditional_str_parser_output',
         'type': 'schema',
       }),
       dict({
@@ -1830,6 +1848,7 @@
           'name': 'StrOutputParser',
         }),
         'id': 8,
+        'name': 'StrOutputParser',
         'type': 'runnable',
       }),
       dict({
@@ -1843,6 +1862,7 @@
           'name': 'XMLOutputParser',
         }),
         'id': 9,
+        'name': 'XMLOutputParser',
         'type': 'runnable',
       }),
     ]),
@@ -1986,6 +2006,7 @@
       dict({
         'data': '__start__',
         'id': '__start__',
+        'name': '__start__',
         'type': 'schema',
       }),
       dict({
@@ -1999,6 +2020,7 @@
           'name': 'ask_question',
         }),
         'id': 'ask_question',
+        'name': 'ask_question',
         'type': 'runnable',
       }),
       dict({
@@ -2012,11 +2034,13 @@
           'name': 'answer_question',
         }),
         'id': 'answer_question',
+        'name': 'answer_question',
         'type': 'runnable',
       }),
       dict({
         'data': '__end__',
         'id': '__end__',
+        'name': '__end__',
         'type': 'schema',
       }),
     ]),

diff --git a/libs/core/tests/unit_tests/runnables/test_graph.py b/libs/core/tests/unit_tests/runnables/test_graph.py
@@ -81,6 +81,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                 "id": 0,
                 "type": "schema",
                 "data": "PromptInput",
+                "name": "PromptInput"
             },
             {
                 "id": 1,
@@ -89,6 +90,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     "id": ["langchain", "prompts", "prompt", "PromptTemplate"],
                     "name": "PromptTemplate",
                 },
+                "name": "PromptTemplate",
             },
             {
                 "id": 2,
@@ -97,6 +99,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     "id": ["langchain_core", "language_models", "fake", "FakeListLLM"],
                     "name": "FakeListLLM",
                 },
+                "name": "FakeListLLM",
                 "metadata": {"key": 2},
             },
             {
@@ -111,11 +114,13 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     ],
                     "name": "CommaSeparatedListOutputParser",
                 },
+                "name": "CommaSeparatedListOutputParser",
             },
             {
                 "id": 4,
                 "type": "schema",
                 "data": "CommaSeparatedListOutputParserOutput",
+                "name": "CommaSeparatedListOutputParserOutput",
             },
         ],
         "edges": [
@@ -136,6 +141,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     "properties": {"name": {"title": "Name", "type": "string"}},
                     "required": ["name"],
                 },
+                "name": "PromptInput",
             },
             {
                 "id": 1,
@@ -144,6 +150,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     "id": ["langchain", "prompts", "prompt", "PromptTemplate"],
                     "name": "PromptTemplate",
                 },
+                "name": "PromptTemplate",
             },
             {
                 "id": 2,
@@ -153,6 +160,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     "name": "FakeListLLM",
                 },
                 "metadata": {"key": 2},
+                "name": "FakeListLLM",
             },
             {
                 "id": 3,
@@ -166,6 +174,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     ],
                     "name": "CommaSeparatedListOutputParser",
                 },
+                "name": "CommaSeparatedListOutputParser",
             },
             {
                 "id": 4,
@@ -175,6 +184,7 @@ def test_graph_sequence(snapshot: SnapshotAssertion) -> None:
                     "title": "CommaSeparatedListOutputParserOutput",
                     "type": "array",
                 },
+                "name": "CommaSeparatedListOutputParserOutput",
             },
         ],
         "edges": [