Skip to content

Commit e3ea365

Browse files
authored
core: use friendlier names for duplicated nodes in mermaid output (#27747)
Thank you for contributing to LangChain! - [x] **PR title**: "core: use friendlier names for duplicated nodes in mermaid output" - **Description:** When generating the Mermaid visualization of a chain, if the chain had multiple nodes of the same type, the reid function would replace their names with the UUID node_id. This made the generated graph difficult to understand. This change deduplicates the nodes in a chain by appending an index to their names. - **Issue:** None - **Discussion:** #27714 - **Dependencies:** None - [ ] **Add tests and docs**: - Currently this functionality is not covered by unit tests, happy to add tests if you'd like - [x] **Lint and test**: Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified. See contribution guidelines for more: https://python.langchain.com/docs/contributing/ Additional guidelines: - Make sure optional dependencies are imported within a function. - Please do not add dependencies to pyproject.toml files (even optional ones) unless they are required for unit tests. - Most PRs should not touch more than one package. - Changes should be backwards compatible. - If you are adding something to community, do not re-import it in langchain. If no one reviews your PR within a few days, please @-mention one of baskaryan, efriis, eyurtsev, ccurme, vbarda, hwchase17. # Example Code: ```python from langchain_core.runnables import RunnablePassthrough def fake_llm(prompt: str) -> str: # Fake LLM for the example return "completion" runnable = { 'llm1': fake_llm, 'llm2': fake_llm, } | RunnablePassthrough.assign( total_chars=lambda inputs: len(inputs['llm1'] + inputs['llm2']) ) print(runnable.get_graph().draw_mermaid(with_styles=False)) ``` # Before ```mermaid graph TD; Parallel_llm1_llm2_Input --> 0b01139db5ed4587ad37964e3a40c0ec; 0b01139db5ed4587ad37964e3a40c0ec --> Parallel_llm1_llm2_Output; Parallel_llm1_llm2_Input --> a98d4b56bd294156a651230b9293347f; a98d4b56bd294156a651230b9293347f --> Parallel_llm1_llm2_Output; Parallel_total_chars_Input --> Lambda; Lambda --> Parallel_total_chars_Output; Parallel_total_chars_Input --> Passthrough; Passthrough --> Parallel_total_chars_Output; Parallel_llm1_llm2_Output --> Parallel_total_chars_Input; ``` # After ```mermaid graph TD; Parallel_llm1_llm2_Input --> fake_llm_1; fake_llm_1 --> Parallel_llm1_llm2_Output; Parallel_llm1_llm2_Input --> fake_llm_2; fake_llm_2 --> Parallel_llm1_llm2_Output; Parallel_total_chars_Input --> Lambda; Lambda --> Parallel_total_chars_Output; Parallel_total_chars_Input --> Passthrough; Passthrough --> Parallel_total_chars_Output; Parallel_llm1_llm2_Output --> Parallel_total_chars_Input; ```
1 parent 71f590d commit e3ea365

File tree

3 files changed

+40
-5
lines changed

3 files changed

+40
-5
lines changed

libs/core/langchain_core/runnables/graph.py

+12-5
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from __future__ import annotations
22

33
import inspect
4-
from collections import Counter
4+
from collections import defaultdict
55
from collections.abc import Sequence
66
from dataclasses import dataclass, field
77
from enum import Enum
@@ -423,12 +423,19 @@ def prefixed(id: str) -> str:
423423
def reid(self) -> Graph:
424424
"""Return a new graph with all nodes re-identified,
425425
using their unique, readable names where possible."""
426-
node_labels = {node.id: node.name for node in self.nodes.values()}
427-
node_label_counts = Counter(node_labels.values())
426+
node_name_to_ids = defaultdict(list)
427+
for node in self.nodes.values():
428+
node_name_to_ids[node.name].append(node.id)
429+
430+
unique_labels = {
431+
node_id: node_name if len(node_ids) == 1 else f"{node_name}_{i + 1}"
432+
for node_name, node_ids in node_name_to_ids.items()
433+
for i, node_id in enumerate(node_ids)
434+
}
428435

429436
def _get_node_id(node_id: str) -> str:
430-
label = node_labels[node_id]
431-
if is_uuid(node_id) and node_label_counts[label] == 1:
437+
label = unique_labels[node_id]
438+
if is_uuid(node_id):
432439
return label
433440
else:
434441
return node_id

libs/core/tests/unit_tests/runnables/__snapshots__/test_graph.ambr

+14
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,20 @@
2626

2727
'''
2828
# ---
29+
# name: test_graph_mermaid_duplicate_nodes[mermaid]
30+
'''
31+
graph TD;
32+
PromptInput --> PromptTemplate_1;
33+
Parallel_llm1_llm2_Input --> FakeListLLM_1;
34+
FakeListLLM_1 --> Parallel_llm1_llm2_Output;
35+
Parallel_llm1_llm2_Input --> FakeListLLM_2;
36+
FakeListLLM_2 --> Parallel_llm1_llm2_Output;
37+
PromptTemplate_1 --> Parallel_llm1_llm2_Input;
38+
PromptTemplate_2 --> PromptTemplateOutput;
39+
Parallel_llm1_llm2_Output --> PromptTemplate_2;
40+
41+
'''
42+
# ---
2943
# name: test_graph_sequence[ascii]
3044
'''
3145
+-------------+

libs/core/tests/unit_tests/runnables/test_graph.py

+14
Original file line numberDiff line numberDiff line change
@@ -405,3 +405,17 @@ def test_graph_mermaid_escape_node_label() -> None:
405405
assert _escape_node_label("foo-bar") == "foo-bar"
406406
assert _escape_node_label("foo_1") == "foo_1"
407407
assert _escape_node_label("#foo*&!") == "_foo___"
408+
409+
410+
def test_graph_mermaid_duplicate_nodes(snapshot: SnapshotAssertion) -> None:
411+
fake_llm = FakeListLLM(responses=["foo", "bar"])
412+
sequence: Runnable = (
413+
PromptTemplate.from_template("Hello, {input}")
414+
| {
415+
"llm1": fake_llm,
416+
"llm2": fake_llm,
417+
}
418+
| PromptTemplate.from_template("{llm1} {llm2}")
419+
)
420+
graph = sequence.get_graph()
421+
assert graph.draw_mermaid(with_styles=False) == snapshot(name="mermaid")

0 commit comments

Comments
 (0)