diff --git a/python/packages/core/agent_framework/_workflows/_viz.py b/python/packages/core/agent_framework/_workflows/_viz.py index e71d917f7aa..55d4f373a84 100644 --- a/python/packages/core/agent_framework/_workflows/_viz.py +++ b/python/packages/core/agent_framework/_workflows/_viz.py @@ -4,10 +4,11 @@ import re import tempfile import uuid +from collections.abc import Iterator from pathlib import Path from typing import Literal -from ._edge import FanInEdgeGroup, InternalEdgeGroup +from ._edge import Edge, FanInEdgeGroup, InternalEdgeGroup from ._workflow import Workflow # Import of WorkflowExecutor is performed lazily inside methods to avoid cycles @@ -189,12 +190,14 @@ def to_mermaid(self, include_internal_executors: bool = False) -> str: A string representation of the workflow in Mermaid flowchart syntax. """ lines: list[str] = ["flowchart TD"] + node_ids = self._mermaid_node_ids(include_internal_executors=include_internal_executors) # Emit top-level workflow self._emit_workflow_mermaid( self._workflow, lines, indent=" ", + node_ids=node_ids, include_internal_executors=include_internal_executors, ) @@ -203,6 +206,7 @@ def to_mermaid(self, include_internal_executors: bool = False) -> str: self._workflow, lines, indent=" ", + node_ids=node_ids, include_internal_executors=include_internal_executors, ) @@ -236,17 +240,26 @@ def _compute_normal_edges( include_internal_executors: bool = False, ) -> list[tuple[str, str, bool]]: """Return list of (source_id, target_id, is_conditional) for non-fan-in groups.""" - edges: list[tuple[str, str, bool]] = [] + return [ + (edge.source_id, edge.target_id, getattr(edge, "_condition", None) is not None) + for edge, _ in self._iter_normal_edges(workflow, include_internal_executors) + ] + + def _iter_normal_edges( + self, + workflow: Workflow | None = None, + include_internal_executors: bool = False, + ) -> Iterator[tuple[Edge, bool]]: + """Retain internal-source identity even when its ID matches a real executor.""" workflow = workflow or self._workflow for group in workflow.edge_groups: if isinstance(group, FanInEdgeGroup): continue - if isinstance(group, InternalEdgeGroup) and not include_internal_executors: + is_internal = isinstance(group, InternalEdgeGroup) + if is_internal and not include_internal_executors: continue for edge in group.edges: - is_cond = getattr(edge, "_condition", None) is not None - edges.append((edge.source_id, edge.target_id, is_cond)) - return edges + yield edge, is_internal # endregion @@ -343,24 +356,64 @@ def _emit_sub_workflows_digraph( # region Internal emitters (Mermaid) + @staticmethod + def _sanitize_mermaid_id(value: str) -> str: + sanitized = re.sub(r"[^0-9A-Za-z_]", "_", value) + return sanitized if sanitized and sanitized[0].isalpha() else f"n_{sanitized}" + + def _mermaid_node_ids(self, *, include_internal_executors: bool = False) -> dict[tuple[str, ...], str]: + """Allocate graph-wide aliases without consuming another node's ordinary ID.""" + from ._workflow_executor import WorkflowExecutor + + candidates: dict[tuple[str, ...], str] = {} + + def collect(workflow: Workflow, ns: tuple[str, ...] = ()) -> None: + prefix = f"{self._sanitize_mermaid_id(ns[-1])}__" if ns else "" + for executor_id, executor in workflow.executors.items(): + candidates[("executor", *ns, executor_id)] = prefix + self._sanitize_mermaid_id(executor_id) + if isinstance(executor, WorkflowExecutor) and executor.workflow: + if ns: + candidates[("subgraph", *ns, executor_id)] = self._sanitize_mermaid_id(executor_id) + collect(executor.workflow, (*ns, executor_id)) + for edge, is_internal in self._iter_normal_edges(workflow, include_internal_executors): + source_kind = "internal" if is_internal else "executor" + candidates[(source_kind, *ns, edge.source_id)] = prefix + self._sanitize_mermaid_id(edge.source_id) + candidates[("executor", *ns, edge.target_id)] = prefix + self._sanitize_mermaid_id(edge.target_id) + for dot_node_id, _, target in self._compute_fan_in_descriptors(workflow): + digest = dot_node_id.split("::")[-1] + candidates[("fan-in", *ns, dot_node_id)] = ( + f"fan_in__{prefix}{self._sanitize_mermaid_id(f'{target}__{digest}')}" + ) + + collect(self._workflow) + reserved = set(candidates.values()) + used: set[str] = set() + node_ids: dict[tuple[str, ...], str] = {} + # Prefer unchanged executor IDs, then allocate deterministically by original scope. + for key in sorted(candidates, key=lambda key: (candidates[key] != key[-1], key)): + candidate = candidates[key] + alias = candidate + if alias in used: + suffix = 2 + alias = f"{candidate}_{suffix}" + while alias in reserved or alias in used: + suffix += 1 + alias = f"{candidate}_{suffix}" + used.add(alias) + node_ids[key] = alias + return node_ids + def _emit_workflow_mermaid( self, workflow: Workflow, lines: list[str], indent: str, - ns: str | None = None, + node_ids: dict[tuple[str, ...], str], + ns: tuple[str, ...] = (), include_internal_executors: bool = False, ) -> None: - def _san(s: str) -> str: - s2 = re.sub(r"[^0-9A-Za-z_]", "_", s) - if not s2 or not s2[0].isalpha(): - s2 = f"n_{s2}" - return s2 - def map_id(x: str) -> str: - if ns: - return f"{_san(ns)}__{_san(x)}" - return _san(x) + return node_ids[("executor", *ns, x)] # Nodes start_executor_id = workflow.start_executor_id @@ -374,9 +427,7 @@ def map_id(x: str) -> str: fan_in_nodes_dot = self._compute_fan_in_descriptors(workflow) fan_in_nodes: list[tuple[str, list[str], str]] = [] for dot_node_id, sources, target in fan_in_nodes_dot: - digest = dot_node_id.split("::")[-1] - base = f"{target}__{digest}" - fan_node_id = f"fan_in__{_san(ns) + '__' if ns else ''}{_san(base)}" + fan_node_id = node_ids[("fan-in", *ns, dot_node_id)] fan_in_nodes.append((fan_node_id, sources, target)) for fan_node_id, _, _ in fan_in_nodes: @@ -390,12 +441,11 @@ def map_id(x: str) -> str: lines.append(f"{indent}{fan_node_id} --> {map_id(target)};") # Normal edges - for src, tgt, is_cond in self._compute_normal_edges( - workflow, include_internal_executors=include_internal_executors - ): - s = map_id(src) - t = map_id(tgt) - if is_cond: + for edge, is_internal in self._iter_normal_edges(workflow, include_internal_executors): + source_kind = "internal" if is_internal else "executor" + s = node_ids[(source_kind, *ns, edge.source_id)] + t = map_id(edge.target_id) + if getattr(edge, "_condition", None) is not None: lines.append(f"{indent}{s} -. conditional .-> {t};") else: lines.append(f"{indent}{s} --> {t};") @@ -405,6 +455,8 @@ def _emit_sub_workflows_mermaid( workflow: Workflow, lines: list[str], indent: str, + node_ids: dict[tuple[str, ...], str], + ns: tuple[str, ...] = (), include_internal_executors: bool = False, ) -> None: try: @@ -412,22 +464,17 @@ def _emit_sub_workflows_mermaid( except ImportError: # pragma: no cover return - def _san(s: str) -> str: - s2 = re.sub(r"[^0-9A-Za-z_]", "_", s) - if not s2 or not s2[0].isalpha(): - s2 = f"n_{s2}" - return s2 - for exec_id, exec_obj in workflow.executors.items(): if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow: - sg_id = _san(exec_id) + sg_id = node_ids[("subgraph", *ns, exec_id)] if ns else node_ids["executor", exec_id] lines.append(f"{indent}subgraph {sg_id}") # Render nested workflow within this subgraph using namespacing self._emit_workflow_mermaid( exec_obj.workflow, lines, indent=f"{indent} ", - ns=exec_id, + node_ids=node_ids, + ns=(*ns, exec_id), include_internal_executors=include_internal_executors, ) # Recurse into deeper sub-workflows @@ -435,6 +482,8 @@ def _san(s: str) -> str: exec_obj.workflow, lines, indent=f"{indent} ", + node_ids=node_ids, + ns=(*ns, exec_id), include_internal_executors=include_internal_executors, ) lines.append(f"{indent}end") diff --git a/python/packages/core/tests/workflow/test_viz.py b/python/packages/core/tests/workflow/test_viz.py index 7b5e42c6cad..049d754c54a 100644 --- a/python/packages/core/tests/workflow/test_viz.py +++ b/python/packages/core/tests/workflow/test_viz.py @@ -2,6 +2,7 @@ """Tests for the workflow visualization module.""" +import re from pathlib import Path from typing import Any @@ -410,3 +411,204 @@ def test_workflow_viz_nested_sub_workflows(): # Should have multiple subgraphs for nested structure subgraph_count = dot_content.count("subgraph cluster_") assert subgraph_count >= 2 # At least one for each level of nesting + + +@pytest.mark.parametrize( + "executor_ids", + [ + ("step-a", "step_a"), + ("step a", "step-a", "step_a", "step_a_2"), + ("步骤甲", "步骤乙"), + ("1step", "n_1step"), + ("_step", "n__step"), + ("step-a", "step_b"), + ("step_a", "step_b"), + ], +) +def test_workflow_viz_mermaid_distinct_executor_ids(executor_ids: tuple[str, ...]) -> None: + """Distinct executors must keep distinct diagram nodes and their original edges.""" + executors = [MockExecutor(id=executor_id) for executor_id in executor_ids] + builder = WorkflowBuilder(start_executor=executors[0]) + for source, target in zip(executors, executors[1:]): + builder.add_edge(source, target) + workflow = builder.build() + viz = WorkflowViz(workflow) + mermaid = viz.to_mermaid() + declarations = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["([^"\n]*)"\];$', mermaid, re.MULTILINE) + node_ids = [node_id for node_id, _ in declarations] + + assert len(node_ids) == len(executor_ids) + assert len(set(node_ids)) == len(executor_ids) + assert [label for _, label in declarations] == [f"{executor_ids[0]} (Start)", *executor_ids[1:]] + for source, target in zip(node_ids, node_ids[1:]): + assert f"{source} --> {target};" in mermaid + assert mermaid == viz.to_mermaid() + # Ordinary identifiers must not be consumed by aliases generated for other nodes. + for executor_id, node_id in zip(executor_ids, node_ids): + if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", executor_id): + assert node_id == executor_id + + +def test_workflow_viz_empty_executor_id_rejected() -> None: + """Visualization does not relax the existing non-empty executor ID contract.""" + with pytest.raises(ValueError, match="Executor ID must be a non-empty string"): + MockExecutor(id="") + + +def test_workflow_viz_mermaid_subgraph_id_collisions() -> None: + """Colliding wrapper aliases must keep separate subgraphs and scoped edges.""" + wrappers = [] + for wrapper_id in ("container-a", "container_a"): + source, target = MockExecutor(id="step-a"), MockExecutor(id="step_a") + sub_workflow = WorkflowBuilder(start_executor=source).add_edge(source, target).build() + wrappers.append(WorkflowExecutor(sub_workflow, id=wrapper_id)) + start = MockExecutor(id="start") + workflow = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + mermaid = WorkflowViz(workflow).to_mermaid() + nodes = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', mermaid, re.MULTILINE) + subgraphs = re.findall(r" subgraph ([A-Za-z][A-Za-z0-9_]*)\n(.*?)\n end", mermaid, re.DOTALL) + + assert len(nodes) == len(set(nodes)) == 7 + assert len(subgraphs) == len({node_id for node_id, _ in subgraphs}) == 2 + for wrapper_id, (_, content) in zip(("container-a", "container_a"), subgraphs): + wrapper_match = re.search(rf'([A-Za-z][A-Za-z0-9_]*)\["{re.escape(wrapper_id)}"\];', mermaid) + assert wrapper_match is not None + child_ids = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', content, re.MULTILINE) + assert len(child_ids) == 2 + assert f"{child_ids[0]} --> {child_ids[1]};" in content + assert f"start --> {wrapper_match.group(1)};" in mermaid + + +def test_workflow_viz_mermaid_namespace_collision() -> None: + """A namespaced child cannot reuse an ordinary top-level executor ID.""" + child = MockExecutor(id="child") + wrapper = WorkflowExecutor(WorkflowBuilder(start_executor=child).build(), id="container") + start = MockExecutor(id="container__child") + workflow = WorkflowBuilder(start_executor=start).add_edge(start, wrapper).build() + mermaid = WorkflowViz(workflow).to_mermaid() + nodes = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', mermaid, re.MULTILINE) + + assert len(nodes) == len(set(nodes)) == 3 + assert 'container__child["container__child (Start)"];' in mermaid + assert "container__child --> container;" in mermaid + + +def test_workflow_viz_mermaid_nested_namespace_collision() -> None: + """Full parent scopes distinguish equal inner wrapper and child names.""" + wrappers = [] + for wrapper_id in ("left", "right"): + leaf = MockExecutor(id="leaf") + inner = WorkflowExecutor(WorkflowBuilder(start_executor=leaf).build(), id="inner") + wrappers.append(WorkflowExecutor(WorkflowBuilder(start_executor=inner).build(), id=wrapper_id)) + start = MockExecutor(id="start") + workflow = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + viz = WorkflowViz(workflow) + mermaid = viz.to_mermaid() + nodes = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', mermaid, re.MULTILINE) + subgraphs = re.findall(r"^\s*subgraph ([A-Za-z][A-Za-z0-9_]*)$", mermaid, re.MULTILINE) + + assert len(nodes) == len(set(nodes)) == 7 + assert len(subgraphs) == len(set(subgraphs)) == 4 + assert mermaid == viz.to_mermaid() + + +def test_workflow_viz_mermaid_fan_in_id_collision() -> None: + """A generated fan-in node cannot reuse a legal executor ID.""" + start, source = MockExecutor(id="start"), MockExecutor(id="source") + target = ListStrTargetExecutor(id="target") + digest = WorkflowViz(WorkflowBuilder(start_executor=start).build())._fan_in_digest("target", ["start", "source"]) + ordinary_id = f"fan_in__target__{digest}" + ordinary = MockExecutor(id=ordinary_id) + workflow = ( + WorkflowBuilder(start_executor=start) + .add_edge(start, source) + .add_edge(start, ordinary) + .add_fan_in_edges([start, source], target) + .build() + ) + mermaid = WorkflowViz(workflow).to_mermaid() + fan_in_ids = re.findall(r"^\s*([A-Za-z][A-Za-z0-9_]*)\(\(fan-in\)\)$", mermaid, re.MULTILINE) + nodes = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', mermaid, re.MULTILINE) + + assert len(fan_in_ids) == 1 + assert len({*nodes, *fan_in_ids}) == 5 + assert f'{ordinary_id}["{ordinary_id}"];' in mermaid + assert f"start --> {fan_in_ids[0]};" in mermaid + assert f"source --> {fan_in_ids[0]};" in mermaid + assert f"{fan_in_ids[0]} --> target;" in mermaid + + +@pytest.mark.parametrize("include_internal_executors", [False, True]) +@pytest.mark.parametrize( + "executor_ids", + [ + ("start", "target", "finish"), + ("foo", "internal:foo", "internal:internal:foo"), + ("step-a", "internal_step_a", "internal_step_a_2"), + ("1step", "internal_1step", "internal_1step_2"), + ], +) +def test_workflow_viz_mermaid_internal_endpoints( + executor_ids: tuple[str, ...], include_internal_executors: bool +) -> None: + """Internal edge sources keep distinct aliases alongside ordinary executor IDs.""" + executors = [MockExecutor(id=executor_id) for executor_id in executor_ids] + builder = WorkflowBuilder(start_executor=executors[0]) + for source, target in zip(executors, executors[1:]): + builder.add_edge(source, target) + viz = WorkflowViz(builder.build()) + mermaid = viz.to_mermaid(include_internal_executors=include_internal_executors) + nodes = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', mermaid, re.MULTILINE) + edges = re.findall(r"^\s*([A-Za-z][A-Za-z0-9_]*) --> ([A-Za-z][A-Za-z0-9_]*);$", mermaid, re.MULTILINE) + ordinary_edges = list(zip(nodes, nodes[1:])) + internal_edges = [(source, target) for source, target in edges if (source, target) not in ordinary_edges] + + assert len(nodes) == len(set(nodes)) == 3 + assert all(edge in edges for edge in ordinary_edges) + if include_internal_executors: + internal_sources = {source for source, _ in internal_edges} + assert len(internal_edges) == len(internal_sources) == 3 + assert internal_sources.isdisjoint(nodes) + assert {target for _, target in internal_edges} == set(nodes) + else: + assert not internal_edges + for executor_id, node_id in zip(executor_ids, nodes): + if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", executor_id): + assert node_id == executor_id + assert mermaid == viz.to_mermaid(include_internal_executors=include_internal_executors) + + +@pytest.mark.parametrize("include_internal_executors", [False, True]) +@pytest.mark.parametrize("colliding_executor", [False, True]) +def test_workflow_viz_mermaid_nested_internal_endpoints( + include_internal_executors: bool, colliding_executor: bool +) -> None: + """Equal internal endpoints in separate nested scopes must remain separate nodes.""" + wrappers = [] + for wrapper_id in ("left", "right"): + leaf = MockExecutor(id="leaf") + builder = WorkflowBuilder(start_executor=leaf) + if colliding_executor: + builder.add_edge(leaf, MockExecutor(id="internal:leaf")) + inner = WorkflowExecutor(builder.build(), id="inner") + wrappers.append(WorkflowExecutor(WorkflowBuilder(start_executor=inner).build(), id=wrapper_id)) + start = MockExecutor(id="inner__internal_leaf") + workflow = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + viz = WorkflowViz(workflow) + mermaid = viz.to_mermaid(include_internal_executors=include_internal_executors) + nodes = re.findall(r'^\s*([A-Za-z][A-Za-z0-9_]*)\["', mermaid, re.MULTILINE) + subgraphs = re.findall(r"^\s*subgraph ([A-Za-z][A-Za-z0-9_]*)$", mermaid, re.MULTILINE) + edges = re.findall(r"^\s*([A-Za-z][A-Za-z0-9_]*) --> ([A-Za-z][A-Za-z0-9_]*);$", mermaid, re.MULTILINE) + internal_edges = [(source, target) for source, target in edges if source not in nodes] + + assert len(nodes) == len(set(nodes)) == (9 if colliding_executor else 7) + assert len(subgraphs) == len(set(subgraphs)) == 4 + assert nodes[0] == "inner__internal_leaf" + if include_internal_executors: + internal_sources = {source for source, _ in internal_edges} + assert len(internal_edges) == len(internal_sources) == (9 if colliding_executor else 7) + assert internal_sources.isdisjoint(subgraphs) + assert {target for _, target in internal_edges} == set(nodes) + else: + assert not internal_edges + assert mermaid == viz.to_mermaid(include_internal_executors=include_internal_executors)