Add resource config node support to workflow representation (#294)

This commit is contained in:
Adrian Lyjak
2026-01-23 16:14:30 -05:00
committed by GitHub
parent be198692de
commit 6dd7fc0de8
9 changed files with 861 additions and 255 deletions
+6
View File
@@ -0,0 +1,6 @@
---
"llama-index-utils-workflow": minor
"llama-index-workflows": minor
---
Add resource config node support to workflow representation
@@ -23,6 +23,7 @@ from workflows.representation import (
WorkflowGraph,
WorkflowGraphEdge,
WorkflowGraphNode,
WorkflowResourceConfigNode,
WorkflowResourceNode,
)
from workflows.representation import (
@@ -45,6 +46,8 @@ def _get_node_color(node: WorkflowGraphNode) -> str:
return "#BEDAE4" # Light blue-gray for external
elif node.node_type == "resource":
return "#DDA0DD" # Plum/light purple for resources
elif node.node_type == "resource_config":
return "#B2DFDB" # Light teal for resource configs
elif node.node_type == "event":
if node.is_subclass_of("StartEvent"):
return "#E27AFF" # Pink for start events
@@ -79,6 +82,8 @@ def _get_node_shape(node: WorkflowGraphNode) -> str:
return "ellipse"
elif node.node_type == "resource":
return "hexagon"
elif node.node_type == "resource_config":
return "box"
elif node.node_type in ("agent", "tool", "workflow_agent", "workflow_handoff"):
return "ellipse"
elif node.node_type == "workflow_base":
@@ -125,6 +130,15 @@ def _render_pyvis(
if node.description:
title_parts.append(f"Doc: {node.description[:100]}...")
title = "\n".join(title_parts)
elif isinstance(node, WorkflowResourceConfigNode):
title_parts = []
if node.type_name:
title_parts.append(f"Type: {node.type_name}")
if node.config_file:
title_parts.append(f"File: {node.config_file}")
if node.path_selector:
title_parts.append(f"Path: {node.path_selector}")
title = "\n".join(title_parts) if title_parts else title
net.add_node(
node.id,
@@ -171,6 +185,8 @@ def _get_mermaid_css_class(node: WorkflowGraphNode) -> str:
return "externalStyle"
elif node.node_type == "resource":
return "resourceStyle"
elif node.node_type == "resource_config":
return "resourceConfigStyle"
elif node.node_type == "event":
if node.is_subclass_of("StartEvent"):
return "startEventStyle"
@@ -274,6 +290,7 @@ def _render_mermaid(
" classDef stepStyle fill:#ADD8E6,color:#000000,line-height:1.2",
" classDef externalStyle fill:#BEDAE4,color:#000000,line-height:1.2",
" classDef resourceStyle fill:#DDA0DD,color:#000000,line-height:1.2",
" classDef resourceConfigStyle fill:#B2DFDB,color:#000000,line-height:1.2",
" classDef startEventStyle fill:#E27AFF,color:#000000",
" classDef stopEventStyle fill:#FFA07A,color:#000000",
" classDef defaultEventStyle fill:#90EE90,color:#000000",
@@ -702,6 +719,7 @@ def draw_most_recent_execution_mermaid(
"classDef stepStyle fill:#ADD8E6,color:#000000,line-height:1.2",
"classDef externalStyle fill:#BEDAE4,color:#000000,line-height:1.2",
"classDef resourceStyle fill:#DDA0DD,color:#000000,line-height:1.2",
"classDef resourceConfigStyle fill:#B2DFDB,color:#000000,line-height:1.2",
"classDef startEventStyle fill:#E27AFF,color:#000000",
"classDef stopEventStyle fill:#FFA07A,color:#000000",
"classDef defaultEventStyle fill:#90EE90,color:#000000",
@@ -1,3 +1,6 @@
import json
from pathlib import Path
from typing import Annotated
from unittest.mock import MagicMock, mock_open, patch
import pytest
@@ -7,6 +10,10 @@ from llama_index.utils.workflow import (
draw_most_recent_execution,
draw_most_recent_execution_mermaid,
)
from pydantic import BaseModel
from workflows.decorators import step
from workflows.events import StartEvent, StopEvent
from workflows.resource import Resource, ResourceConfig
from workflows.workflow import Workflow
@@ -430,6 +437,107 @@ def test_mermaid_resource_style_always_defined(workflow: Workflow) -> None:
assert "classDef resourceStyle fill:#DDA0DD" in result
def test_mermaid_resource_config_nodes_rendered(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Test that resource config nodes are rendered in Mermaid output."""
monkeypatch.chdir(tmp_path)
config_data = {"setting": "test", "value": 1}
with open("config.json", "w") as f:
json.dump(config_data, f)
class ConfigData(BaseModel):
setting: str
value: int
class Client:
pass
def get_client(
config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> Client:
return Client()
class WorkflowWithConfig(Workflow):
@step
async def start_step(
self,
ev: StartEvent,
client: Annotated[Client, Resource(get_client)],
) -> StopEvent:
return StopEvent(result="done")
result = draw_all_possible_flows_mermaid(WorkflowWithConfig())
assert "classDef resourceConfigStyle fill:#B2DFDB" in result
lines = result.split("\n")
resource_config_lines = [
line
for line in lines
if "resource_config_" in line and ":::" in line and " --> " not in line
]
assert len(resource_config_lines) == 1
assert ":::resourceConfigStyle" in resource_config_lines[0]
def test_pyvis_resource_config_nodes_rendered(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Test that resource config nodes are rendered in Pyvis output."""
monkeypatch.chdir(tmp_path)
config_data = {"setting": "test", "value": 1}
with open("config.json", "w") as f:
json.dump(config_data, f)
class ConfigData(BaseModel):
setting: str
value: int
class Client:
pass
def get_client(
config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> Client:
return Client()
class WorkflowWithConfig(Workflow):
@step
async def start_step(
self,
ev: StartEvent,
client: Annotated[Client, Resource(get_client)],
) -> StopEvent:
return StopEvent(result="done")
with patch("llama_index.utils.workflow.Network") as mock_network:
mock_net_instance = MagicMock()
mock_network.return_value = mock_net_instance
draw_all_possible_flows(WorkflowWithConfig(), filename="test.html")
node_calls = mock_net_instance.add_node.call_args_list
resource_config_nodes = []
for call in node_calls:
args, kwargs = call
node_id = args[0]
if "resource_config_" in node_id:
resource_config_nodes.append((node_id, kwargs))
assert len(resource_config_nodes) == 1, "Should have resource config node"
node_id, kwargs = resource_config_nodes[0]
assert kwargs.get("shape") == "box", (
f"Resource config node {node_id} should be box"
)
assert kwargs.get("color") == "#B2DFDB", (
f"Resource config node {node_id} should be light teal"
)
def test_resource_node_deduplication_in_rendering(
workflow_with_resources: Workflow,
) -> None:
@@ -7,6 +7,7 @@ from workflows.representation.types import (
WorkflowGraphEdge,
WorkflowGraphNode,
WorkflowNodeBase,
WorkflowResourceConfigNode,
WorkflowResourceNode,
WorkflowStepNode,
)
@@ -18,6 +19,7 @@ __all__ = [
"WorkflowEventNode",
"WorkflowExternalNode",
"WorkflowResourceNode",
"WorkflowResourceConfigNode",
"WorkflowGenericNode",
"WorkflowGraphNode",
"WorkflowGraphEdge",
@@ -2,9 +2,12 @@ from __future__ import annotations
import hashlib
import inspect
from typing import Any
from pydantic import BaseModel
from workflows import Workflow
from workflows.decorators import StepConfig, StepFunction
from workflows.decorators import StepFunction
from workflows.events import (
Event,
HumanResponseEvent,
@@ -17,12 +20,14 @@ from workflows.representation.types import (
WorkflowGraph,
WorkflowGraphEdge,
WorkflowGraphNode,
WorkflowResourceConfigNode,
WorkflowResourceNode,
WorkflowStepNode,
)
from workflows.resource import (
ResourceDefinition,
ResourceDescriptor,
_get_resource_config_data,
_Resource,
_ResourceConfig,
)
@@ -32,6 +37,15 @@ from workflows.utils import (
)
def _get_type_name(type_annotation: type | None) -> str | None:
"""Extract a readable type name from a type annotation."""
if type_annotation is None:
return None
if hasattr(type_annotation, "__name__"):
return type_annotation.__name__
return str(type_annotation)
def _get_resource_identity(resource: ResourceDescriptor) -> int:
"""Get a unique identifier for resource deduplication.
@@ -62,6 +76,47 @@ def _get_event_type_chain(cls: type) -> list[str]:
return names
def _create_resource_config_node(
resource_config: _ResourceConfig,
type_annotation: type | None,
) -> WorkflowResourceConfigNode:
"""Create a WorkflowResourceConfigNode from a _ResourceConfig."""
# Compute unique hash for deduplication based on config file and path selector
hash_input = f"{resource_config.config_file}:{resource_config.path_selector or ''}"
unique_hash = hashlib.sha256(hash_input.encode()).hexdigest()[:12]
node_id = f"resource_config_{unique_hash}"
type_name = _get_type_name(type_annotation)
# Prefer explicit label, then type name, then config file path
label = resource_config.label or type_name or resource_config.config_file
# Extract JSON schema if type is a BaseModel
config_schema: dict[str, Any] | None = None
if (
type_annotation is not None
and isinstance(type_annotation, type)
and issubclass(type_annotation, BaseModel)
):
model_cls: type[BaseModel] = type_annotation
config_schema = model_cls.model_json_schema()
# Read config value using existing infrastructure
config_value = _get_resource_config_data(
resource_config.config_file, resource_config.path_selector
)
return WorkflowResourceConfigNode(
id=node_id,
label=label,
type_name=type_name,
config_file=resource_config.config_file,
path_selector=resource_config.path_selector,
config_schema=config_schema,
config_value=config_value,
description=resource_config.description,
)
def _create_resource_node(resource_def: ResourceDefinition) -> WorkflowResourceNode:
"""Create a WorkflowResourceNode from a ResourceDefinition.
@@ -69,15 +124,7 @@ def _create_resource_node(resource_def: ResourceDefinition) -> WorkflowResourceN
rather than at Resource creation time for performance.
"""
resource = resource_def.resource
# Get type name from annotation
type_name: str | None = None
if resource_def.type_annotation is not None:
type_annotation = resource_def.type_annotation
if hasattr(type_annotation, "__name__"):
type_name = type_annotation.__name__
else:
type_name = str(type_annotation)
type_name = _get_type_name(resource_def.type_annotation)
# Extract source metadata lazily - only available for _Resource with factory
source_file: str | None = None
@@ -86,18 +133,9 @@ def _create_resource_node(resource_def: ResourceDefinition) -> WorkflowResourceN
if isinstance(resource, _Resource):
factory = resource._factory
try:
source_file = inspect.getfile(factory) # type: ignore[arg-type]
except (TypeError, OSError):
pass
try:
_, source_line = inspect.getsourcelines(factory) # type: ignore[arg-type]
except (TypeError, OSError):
pass
source_file = inspect.getfile(factory) # type: ignore[arg-type]
_, source_line = inspect.getsourcelines(factory) # type: ignore[arg-type]
resource_description = inspect.getdoc(factory)
elif isinstance(resource, _ResourceConfig):
# For ResourceConfig, the source is the config file
source_file = resource.config_file
# Compute unique hash for deduplication
hash_input = f"{resource.name}:{source_file or 'unknown'}"
@@ -139,21 +177,102 @@ def get_workflow_representation(workflow: Workflow) -> WorkflowGraph:
nodes: list[WorkflowGraphNode] = []
edges: list[WorkflowGraphEdge] = []
added_nodes: set[str] = set() # Track added node IDs to avoid duplicates
added_resource_nodes: dict[int, WorkflowResourceNode] = {} # Track by factory id
# Track resource nodes by identity (factory id for _Resource)
added_resource_nodes: dict[int, WorkflowResourceNode] = {}
# Track resource config nodes by config_file>path_selector
added_resource_config_nodes: dict[str, WorkflowResourceConfigNode] = {}
# Track descriptor nodes by identity for step edges (_Resource or _ResourceConfig)
added_descriptor_nodes: dict[
int, WorkflowResourceNode | WorkflowResourceConfigNode
] = {}
# Track which resources have had their dependencies expanded
expanded_resources: set[int] = set()
expanding_resources: set[int] = set()
# Track resource dependency edges to avoid duplicates
resource_edge_keys: set[tuple[str, str, str | None]] = set()
step_config: StepConfig | None = None
def _ensure_resource_config_node(
resource_config: _ResourceConfig,
type_annotation: type | None,
) -> WorkflowResourceConfigNode:
selector = resource_config.path_selector
config_key = resource_config.config_file + (
(">" + selector) if selector else ""
)
if config_key in added_resource_config_nodes:
return added_resource_config_nodes[config_key]
node = _create_resource_config_node(resource_config, type_annotation)
nodes.append(node)
added_resource_config_nodes[config_key] = node
return node
def _ensure_resource_node(
resource: _Resource,
type_annotation: type | None,
param_name: str,
) -> WorkflowResourceNode:
resource_id = _get_resource_identity(resource)
if resource_id in added_resource_nodes:
return added_resource_nodes[resource_id]
node = _create_resource_node(
ResourceDefinition(
name=param_name,
resource=resource,
type_annotation=type_annotation,
)
)
nodes.append(node)
added_resource_nodes[resource_id] = node
return node
def _track_resource_edge(
source: str,
target: str,
label: str | None,
) -> None:
key = (source, target, label)
if key in resource_edge_keys:
return
resource_edge_keys.add(key)
edges.append(WorkflowGraphEdge(source=source, target=target, label=label))
def _ensure_descriptor_node(
descriptor: ResourceDescriptor,
type_annotation: type | None,
param_name: str,
) -> WorkflowResourceNode | WorkflowResourceConfigNode:
descriptor_id = _get_resource_identity(descriptor)
if isinstance(descriptor, _ResourceConfig):
node = _ensure_resource_config_node(descriptor, type_annotation)
added_descriptor_nodes[descriptor_id] = node
return node
if isinstance(descriptor, _Resource):
node = _ensure_resource_node(descriptor, type_annotation, param_name)
added_descriptor_nodes[descriptor_id] = node
resource_id = _get_resource_identity(descriptor)
if resource_id in expanded_resources:
return node
if resource_id in expanding_resources:
return node
expanding_resources.add(resource_id)
for dep_name, dep_descriptor, dep_type in descriptor.get_dependencies():
dep_node = _ensure_descriptor_node(dep_descriptor, dep_type, dep_name)
_track_resource_edge(source=node.id, target=dep_node.id, label=dep_name)
expanding_resources.remove(resource_id)
expanded_resources.add(resource_id)
return node
raise TypeError(
f"Unsupported resource descriptor type: {type(descriptor).__name__}"
)
# Only one kind of `StopEvent` is allowed in a `Workflow`.
# Assuming that `Workflow` is validated before drawing, it's enough to find the first one.
current_stop_event = None
for step_name, step_func in steps.items():
step_config = step_func._step_config
for return_type in step_config.return_types:
for step_func in steps.values():
for return_type in step_func._step_config.return_types:
if issubclass(return_type, StopEvent):
current_stop_event = return_type
break
if current_stop_event:
break
@@ -217,11 +336,11 @@ def get_workflow_representation(workflow: Workflow) -> WorkflowGraph:
# Add resource nodes (deduplicated by resource identity)
for resource_def in step_config.resources:
resource_id = _get_resource_identity(resource_def.resource)
if resource_id not in added_resource_nodes:
resource_node = _create_resource_node(resource_def)
nodes.append(resource_node)
added_resource_nodes[resource_id] = resource_node
_ensure_descriptor_node(
resource_def.resource,
resource_def.type_annotation,
resource_def.name,
)
# Second pass: Add edges
for step_name, step_func in steps.items():
@@ -265,7 +384,7 @@ def get_workflow_representation(workflow: Workflow) -> WorkflowGraph:
# Edges from steps to resources (with variable name as label)
for resource_def in step_config.resources:
resource_id = _get_resource_identity(resource_def.resource)
resource_node = added_resource_nodes[resource_id]
resource_node = added_descriptor_nodes[resource_id]
edges.append(
WorkflowGraphEdge(
source=step_name,
@@ -88,6 +88,38 @@ class WorkflowResourceNode(WorkflowNodeBase):
)
class WorkflowResourceConfigNode(WorkflowNodeBase):
"""A resource config node representing a configuration loaded from a JSON file."""
node_type: Literal["resource_config"] = Field(
default="resource_config", description="Discriminator field for node type"
)
type_name: str | None = Field(
default=None,
description="The Pydantic BaseModel type that the config is validated against",
)
config_file: str | None = Field(
default=None,
description="Path to the JSON configuration file",
)
path_selector: str | None = Field(
default=None,
description="Dot-separated path selector for nested configuration values",
)
config_schema: dict[str, Any] | None = Field(
default=None,
description="Pydantic JSON schema for the config type",
)
config_value: dict[str, Any] | None = Field(
default=None,
description="The configuration value read from the file (if readable)",
)
description: str | None = Field(
default=None,
description="Human-readable description of the config's purpose and contents",
)
class WorkflowGenericNode(WorkflowNodeBase):
"""A generic node for custom visualization types not covered by standard node types.
@@ -121,6 +153,7 @@ WorkflowGraphNode = Union[
WorkflowEventNode,
WorkflowExternalNode,
WorkflowResourceNode,
WorkflowResourceConfigNode,
WorkflowGenericNode,
]
@@ -247,6 +280,7 @@ __all__ = [
"WorkflowEventNode",
"WorkflowExternalNode",
"WorkflowResourceNode",
"WorkflowResourceConfigNode",
"WorkflowGenericNode",
"WorkflowGraphNode",
"WorkflowGraphEdge",
@@ -192,12 +192,16 @@ class _ResourceConfig(Generic[B]):
config_file: str
path_selector: str | None
cls_factory: Type[B] | None
label: str | None
description: str | None
def __init__(
self,
config_file: str,
path_selector: str | None,
cls_factory: Type[B] | None = None,
label: str | None = None,
description: str | None = None,
) -> None:
config_path = Path(config_file)
if not config_path.is_file():
@@ -210,6 +214,8 @@ class _ResourceConfig(Generic[B]):
self.config_file = str(config_path.resolve())
self.path_selector = path_selector
self.cls_factory = cls_factory
self.label = label
self.description = description
@property
def name(self) -> str:
@@ -259,20 +265,27 @@ class _ResourceConfig(Generic[B]):
def ResourceConfig(
config_file: str,
path_selector: str | None = None,
label: str | None = None,
description: str | None = None,
) -> _ResourceConfig:
"""
Wrapper for a _ResourceConfig.
Attributes:
config_file (str): JSON file where the configuration is stored
path_selector (str | None): Path selector to retrieve a specific value from the JSON map
cache (bool): Cache the resource's value to avoid re-computation.
Args:
config_file: JSON file where the configuration is stored.
path_selector: Path selector to retrieve a specific value from the JSON map.
label: Human-friendly short name for display in visualizations.
description: Longer description explaining the purpose and contents of this config.
Returns:
_ResourceConfig: A configured resource representation
_ResourceConfig: A configured resource representation.
"""
return _ResourceConfig(config_file=config_file, path_selector=path_selector)
return _ResourceConfig(
config_file=config_file,
path_selector=path_selector,
label=label,
description=description,
)
class ResourceDefinition(BaseModel):
@@ -12,6 +12,7 @@ from pydantic import BaseModel
from workflows.decorators import step
from workflows.errors import WorkflowValidationError
from workflows.events import StartEvent, StopEvent
from workflows.representation import get_workflow_representation
from workflows.resource import Resource, ResourceConfig
from workflows.workflow import Workflow
@@ -142,6 +143,44 @@ async def test_nested_resource_in_resource_with_future_annotations() -> None:
assert result == "done"
def test_resource_config_representation_with_future_annotations(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Representation should include ResourceConfig under future annotations."""
monkeypatch.chdir(tmp_path)
class SimpleConfig(BaseModel):
name: str
class SimpleClient:
def __init__(self, config: SimpleConfig) -> None:
self.config = config
config_path = tmp_path / "config.json"
config_path.write_text(json.dumps({"name": "demo"}))
def get_client(
config: Annotated[SimpleConfig, ResourceConfig(config_file=str(config_path))],
) -> SimpleClient:
return SimpleClient(config=config)
class WorkflowWithConfig(Workflow):
@step
async def start_step(
self,
ev: StartEvent,
client: Annotated[SimpleClient, Resource(get_client)],
) -> StopEvent:
return StopEvent(result="done")
graph = get_workflow_representation(WorkflowWithConfig())
resource_config_nodes = [
node for node in graph.nodes if node.node_type == "resource_config"
]
assert len(resource_config_nodes) == 1
@pytest.mark.asyncio
async def test_localns_does_not_shadow_factory_module_types(
tmp_path: Path,
@@ -1,6 +1,10 @@
from typing import Annotated
import json
from collections.abc import Iterable, Mapping
from pathlib import Path
from typing import Annotated, Optional
import pytest
from pydantic import BaseModel
from workflows.decorators import step
from workflows.events import Event, StartEvent, StopEvent
from workflows.representation import (
@@ -8,16 +12,53 @@ from workflows.representation import (
WorkflowExternalNode,
WorkflowGraph,
WorkflowGraphEdge,
WorkflowGraphNode,
WorkflowResourceConfigNode,
WorkflowResourceNode,
WorkflowStepNode,
get_workflow_representation,
)
from workflows.resource import Resource
from workflows.resource import Resource, ResourceConfig
from workflows.workflow import Workflow
from .conftest import DummyWorkflow # type: ignore[import]
def _nodes_of_type(graph: WorkflowGraph, node_type: str) -> list[WorkflowGraphNode]:
return [node for node in graph.nodes if node.node_type == node_type]
def _resource_nodes(graph: WorkflowGraph) -> list[WorkflowResourceNode]:
return [node for node in graph.nodes if isinstance(node, WorkflowResourceNode)]
def _resource_config_nodes(graph: WorkflowGraph) -> list[WorkflowResourceConfigNode]:
return [
node for node in graph.nodes if isinstance(node, WorkflowResourceConfigNode)
]
def _edges_as_tuples(graph: WorkflowGraph) -> set[tuple[str, str, Optional[str]]]:
return {(edge.source, edge.target, edge.label) for edge in graph.edges}
def _find_edges(
graph: WorkflowGraph,
*,
source: Optional[str] = None,
target_prefix: Optional[str] = None,
label: Optional[str] = None,
) -> list[WorkflowGraphEdge]:
edges: Iterable[WorkflowGraphEdge] = graph.edges
if source is not None:
edges = [edge for edge in edges if edge.source == source]
if target_prefix is not None:
edges = [edge for edge in edges if edge.target.startswith(target_prefix)]
if label is not None:
edges = [edge for edge in edges if edge.label == label]
return list(edges)
@pytest.fixture()
def ground_truth_repr() -> WorkflowGraph:
return WorkflowGraph(
@@ -76,14 +117,12 @@ def test_get_workflow_representation(ground_truth_repr: WorkflowGraph) -> None:
graph = get_workflow_representation(workflow=wf)
assert isinstance(graph, WorkflowGraph)
assert sorted(
[node.id for node in ground_truth_repr.nodes if node.node_type == "step"]
) == sorted([node.id for node in graph.nodes if node.node_type == "step"])
node.id for node in _nodes_of_type(ground_truth_repr, "step")
) == sorted(node.id for node in _nodes_of_type(graph, "step"))
assert sorted(
[node.id for node in ground_truth_repr.nodes if node.node_type == "event"]
) == sorted([node.id for node in graph.nodes if node.node_type == "event"])
expected_edges = ground_truth_repr.edges
for edge in expected_edges:
assert edge in graph.edges
node.id for node in _nodes_of_type(ground_truth_repr, "event")
) == sorted(node.id for node in _nodes_of_type(graph, "event"))
assert _edges_as_tuples(graph) >= _edges_as_tuples(ground_truth_repr)
def test_truncated_label() -> None:
@@ -95,7 +134,7 @@ def test_truncated_label() -> None:
def test_graph_serialization() -> None:
"""Test that WorkflowGraphNodeEdges serializes correctly to JSON."""
"""Test that WorkflowGraph serializes and restores node types."""
graph = WorkflowGraph(
name="TestWorkflow",
nodes=[
@@ -109,36 +148,16 @@ def test_graph_serialization() -> None:
],
edges=[WorkflowGraphEdge(source="test", target="OneTestEvent")],
)
# Test direct access
assert len(graph.nodes) == 2
step_node = graph.nodes[0]
assert isinstance(step_node, WorkflowStepNode)
assert step_node.node_type == "step"
assert step_node.label == "test"
assert step_node.id == "test"
event_node = graph.nodes[1]
assert isinstance(event_node, WorkflowEventNode)
assert event_node.event_type == "OneTestEvent"
assert event_node.event_types == ["OneTestEvent"]
assert event_node.node_type == "event"
assert event_node.label == "OneTestEvent"
assert event_node.id == "OneTestEvent"
assert len(graph.edges) == 1
assert graph.edges[0].source == "test"
assert graph.edges[0].target == "OneTestEvent"
# Test JSON serialization (round-trip works)
data = graph.model_dump()
assert "event_type" not in data["nodes"][0] # Step nodes don't have event_type
assert data["nodes"][1]["event_type"] == "OneTestEvent"
assert data["nodes"][1]["event_types"] == ["OneTestEvent"]
# Test deserialization
restored = WorkflowGraph.model_validate(data)
restored_event = restored.nodes[1]
assert isinstance(restored_event, WorkflowEventNode)
assert restored_event.event_type == "OneTestEvent"
assert restored_event.is_subclass_of("OneTestEvent")
assert len(restored.nodes) == 2
event_node = next(
node for node in restored.nodes if isinstance(node, WorkflowEventNode)
)
assert event_node.event_type == "OneTestEvent"
assert event_node.is_subclass_of("OneTestEvent")
# --- Resource node tests ---
@@ -177,36 +196,25 @@ class WorkflowWithResources(Workflow):
def test_get_workflow_representation_with_resources() -> None:
"""Test that resource nodes are extracted from workflow with resources."""
"""Resource node metadata and step -> resource edge label are derived from factory."""
wf = WorkflowWithResources()
graph = get_workflow_representation(workflow=wf)
# Should have resource nodes
resource_nodes = [n for n in graph.nodes if isinstance(n, WorkflowResourceNode)]
resource_nodes = _resource_nodes(graph)
assert len(resource_nodes) == 1
resource_node = resource_nodes[0]
assert resource_node.node_type == "resource"
assert resource_node.type_name == "DatabaseClient"
assert resource_node.getter_name == "get_database_client"
assert resource_node.description is not None
assert "Factory function" in resource_node.description
assert resource_node.source_file is not None
assert resource_node.source_line is not None
def test_resource_node_edges_have_variable_names() -> None:
"""Test that edges from steps to resources have the variable name as label."""
wf = WorkflowWithResources()
graph = get_workflow_representation(workflow=wf)
# Find edges to resource nodes
resource_edges = [e for e in graph.edges if e.target.startswith("resource_")]
assert len(resource_edges) == 1
edge = resource_edges[0]
assert edge.label == "db_client" # The variable name
assert edge.source == "step_with_resource"
edges = _find_edges(
graph,
source="step_with_resource",
target_prefix="resource_",
label="db_client",
)
assert len(edges) == 1
def test_resource_nodes_are_deduplicated() -> None:
@@ -240,17 +248,12 @@ def test_resource_nodes_are_deduplicated() -> None:
graph = get_workflow_representation(workflow=wf)
# Should have only one resource node (deduplicated)
resource_nodes = [n for n in graph.nodes if isinstance(n, WorkflowResourceNode)]
assert len(resource_nodes) == 1
assert len(_resource_nodes(graph)) == 1
# But should have two edges (one from each step)
resource_edges = [e for e in graph.edges if e.target.startswith("resource_")]
resource_edges = _find_edges(graph, target_prefix="resource_", label="db")
assert len(resource_edges) == 2
# Both edges should have the variable name "db"
for edge in resource_edges:
assert edge.label == "db"
def test_multiple_different_resources() -> None:
"""Test workflow with multiple different resources."""
@@ -275,74 +278,20 @@ def test_multiple_different_resources() -> None:
graph = get_workflow_representation(workflow=wf)
# Should have two different resource nodes
resource_nodes = [n for n in graph.nodes if isinstance(n, WorkflowResourceNode)]
resource_nodes = _resource_nodes(graph)
assert len(resource_nodes) == 2
type_names = {rn.type_name for rn in resource_nodes}
assert type_names == {"DatabaseClient", "CacheClient"}
# Should have two edges with different labels
resource_edges = [e for e in graph.edges if e.target.startswith("resource_")]
resource_edges = _find_edges(graph, target_prefix="resource_")
assert len(resource_edges) == 2
labels = {e.label for e in resource_edges}
assert labels == {"db", "cache"}
def test_resource_node_serialization() -> None:
"""Test that WorkflowResourceNode serializes correctly."""
resource_node = WorkflowResourceNode(
id="resource_abc123",
label="TestType",
type_name="TestType",
getter_name="get_test_type",
source_file="/path/to/file.py",
source_line=42,
description="Test docstring",
)
assert resource_node.id == "resource_abc123"
assert resource_node.label == "TestType"
assert resource_node.node_type == "resource"
assert resource_node.type_name == "TestType"
assert resource_node.getter_name == "get_test_type"
assert resource_node.source_file == "/path/to/file.py"
assert resource_node.source_line == 42
assert resource_node.description == "Test docstring"
# Test serialization
data = resource_node.model_dump()
assert data["id"] == "resource_abc123"
assert data["label"] == "TestType"
assert data["type_name"] == "TestType"
assert data["node_type"] == "resource"
# Test deserialization
restored = WorkflowResourceNode.model_validate(data)
assert restored.id == "resource_abc123"
assert restored.label == "TestType"
assert restored.type_name == "TestType"
assert restored.node_type == "resource"
def test_graph_with_resources() -> None:
"""Test that workflow graph with resources is correct."""
wf = WorkflowWithResources()
graph = get_workflow_representation(workflow=wf)
# Check resource nodes are in the nodes list
resource_nodes = [n for n in graph.nodes if isinstance(n, WorkflowResourceNode)]
assert len(resource_nodes) == 1
rn = resource_nodes[0]
assert rn.type_name == "DatabaseClient"
assert rn.getter_name == "get_database_client"
# Check edges with labels
resource_edges = [e for e in graph.edges if e.label is not None]
assert len(resource_edges) == 1
assert resource_edges[0].label == "db_client"
def test_edge_with_label() -> None:
"""Test that WorkflowGraphEdge with label works correctly."""
edge = WorkflowGraphEdge(source="resource_123", target="my_step", label="my_var")
@@ -361,95 +310,6 @@ def test_edge_without_label() -> None:
assert edge.label is None
# --- Serialization/Deserialization tests for all node types ---
def test_step_node_serialization_roundtrip() -> None:
"""Test WorkflowStepNode serialization and deserialization."""
node = WorkflowStepNode(id="my_step", label="My Step")
data = node.model_dump()
assert data["id"] == "my_step"
assert data["label"] == "My Step"
assert data["node_type"] == "step"
restored = WorkflowStepNode.model_validate(data)
assert restored.id == "my_step"
assert restored.label == "My Step"
assert restored.node_type == "step"
def test_event_node_serialization_roundtrip() -> None:
"""Test WorkflowEventNode serialization and deserialization."""
node = WorkflowEventNode(
id="MyEvent",
label="My Event",
event_type="MyEvent",
event_types=["MyEvent", "ParentEvent"],
)
data = node.model_dump()
assert data["id"] == "MyEvent"
assert data["label"] == "My Event"
assert data["node_type"] == "event"
assert data["event_type"] == "MyEvent"
assert data["event_types"] == ["MyEvent", "ParentEvent"]
restored = WorkflowEventNode.model_validate(data)
assert restored.id == "MyEvent"
assert restored.label == "My Event"
assert restored.node_type == "event"
assert restored.event_type == "MyEvent"
assert restored.event_types == ["MyEvent", "ParentEvent"]
assert restored.is_subclass_of("ParentEvent")
assert not restored.is_subclass_of("UnrelatedEvent")
def test_external_node_serialization_roundtrip() -> None:
"""Test WorkflowExternalNode serialization and deserialization."""
node = WorkflowExternalNode(id="external_step", label="External Step")
data = node.model_dump()
assert data["id"] == "external_step"
assert data["label"] == "External Step"
assert data["node_type"] == "external"
restored = WorkflowExternalNode.model_validate(data)
assert restored.id == "external_step"
assert restored.label == "External Step"
assert restored.node_type == "external"
def test_resource_node_serialization_roundtrip() -> None:
"""Test WorkflowResourceNode serialization and deserialization."""
node = WorkflowResourceNode(
id="resource_abc123",
label="MyResourceType",
type_name="MyResourceType",
getter_name="get_my_resource",
source_file="/path/to/source.py",
source_line=100,
description="Resource docstring",
)
data = node.model_dump()
assert data["id"] == "resource_abc123"
assert data["label"] == "MyResourceType"
assert data["node_type"] == "resource"
assert data["type_name"] == "MyResourceType"
assert data["getter_name"] == "get_my_resource"
assert data["source_file"] == "/path/to/source.py"
assert data["source_line"] == 100
assert data["description"] == "Resource docstring"
restored = WorkflowResourceNode.model_validate(data)
assert restored.id == "resource_abc123"
assert restored.label == "MyResourceType"
assert restored.type_name == "MyResourceType"
assert restored.getter_name == "get_my_resource"
assert restored.node_type == "resource"
def test_graph_with_all_node_types_serialization() -> None:
"""Test full graph serialization/deserialization with all node types."""
graph = WorkflowGraph(
@@ -529,17 +389,28 @@ def test_graph_deserialization_from_raw_json() -> None:
"node_type": "resource",
"type_name": "SomeType",
},
{
"id": "resource_config_456",
"label": "ConfigModel",
"node_type": "resource_config",
"type_name": "ConfigModel",
"config_file": "config.json",
"path_selector": "settings",
"config_schema": {
"type": "object",
"properties": {"key": {"type": "string"}},
},
"config_value": {"key": "value"},
},
],
"edges": [{"source": "MyEvent", "target": "step1"}],
}
graph = WorkflowGraph.model_validate(raw_data)
assert len(graph.nodes) == 4
assert isinstance(graph.nodes[0], WorkflowStepNode)
assert isinstance(graph.nodes[1], WorkflowEventNode)
assert isinstance(graph.nodes[2], WorkflowExternalNode)
assert isinstance(graph.nodes[3], WorkflowResourceNode)
assert len(graph.nodes) == 5
node_types = {node.node_type for node in graph.nodes}
assert node_types == {"step", "event", "external", "resource", "resource_config"}
# --- filter_by_node_type tests ---
@@ -821,3 +692,399 @@ def test_filter_by_node_type_deduplicates_edges() -> None:
assert len(filtered.edges) == 1
assert filtered.edges[0].source == "step1"
assert filtered.edges[0].target == "step2"
# --- Resource config node tests ---
class ConfigData(BaseModel):
"""A config model for testing resource configs."""
setting: str
value: int
def _write_config(tmp_path: Path, filename: str, data: Mapping[str, object]) -> str:
config_path = tmp_path / filename
with open(config_path, "w") as f:
json.dump(data, f)
return str(config_path)
def test_resource_config_nested_in_resource_factory(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Nested ResourceConfig should create resource + config nodes and an edge."""
monkeypatch.chdir(tmp_path)
config_data = {"setting": "test", "value": 42}
config_path = _write_config(tmp_path, "config.json", config_data)
def get_configured_client(
my_config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> DatabaseClient:
return DatabaseClient()
class WorkflowWithResourceConfig(Workflow):
@step
async def step_with_config(
self,
ev: StartEvent,
client: Annotated[DatabaseClient, Resource(get_configured_client)],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithResourceConfig()
graph = get_workflow_representation(workflow=wf)
assert len(_resource_nodes(graph)) == 1
resource_config_nodes = _resource_config_nodes(graph)
assert len(resource_config_nodes) == 1
config_node = resource_config_nodes[0]
assert config_node.type_name == "ConfigData"
assert config_node.config_file == config_path
assert config_node.path_selector is None
assert config_node.config_schema is not None
assert {"setting", "value"} <= set(config_node.config_schema.get("properties", {}))
assert config_node.config_value == config_data
edges = _find_edges(graph, target_prefix="resource_config_", label="my_config")
assert len(edges) == 1
assert edges[0].source.startswith("resource_")
def test_recursive_resource_dependencies_with_config(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Nested resources should create resource->resource and resource->config edges."""
monkeypatch.chdir(tmp_path)
config_data = {"setting": "nested", "value": 7}
config_path = _write_config(tmp_path, "config.json", config_data)
class DBConnection:
def __init__(self, config: ConfigData) -> None:
self.config = config
class Repository:
def __init__(self, db: DBConnection) -> None:
self.db = db
def get_db_connection(
config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> DBConnection:
return DBConnection(config=config)
def get_repository(
db: Annotated[DBConnection, Resource(get_db_connection)],
) -> Repository:
return Repository(db=db)
class WorkflowWithRecursiveResources(Workflow):
@step
async def start_step(
self,
ev: StartEvent,
repo: Annotated[Repository, Resource(get_repository)],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithRecursiveResources()
graph = get_workflow_representation(workflow=wf)
assert len(_resource_nodes(graph)) == 2
resource_config_nodes = _resource_config_nodes(graph)
assert len(resource_config_nodes) == 1
assert resource_config_nodes[0].config_file == config_path
assert resource_config_nodes[0].config_value == config_data
def _resource_node_for_getter(suffix: str) -> WorkflowResourceNode:
return next(
node
for node in _resource_nodes(graph)
if node.getter_name is not None and node.getter_name.endswith(suffix)
)
repo_node = _resource_node_for_getter("get_repository")
db_node = _resource_node_for_getter("get_db_connection")
repo_edges = _find_edges(graph, source=repo_node.id, label="db")
assert len(repo_edges) == 1
assert repo_edges[0].target == db_node.id
config_edges = _find_edges(graph, source=db_node.id, label="config")
assert len(config_edges) == 1
assert config_edges[0].target.startswith("resource_config_")
def test_resource_config_direct_in_step(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""ResourceConfig used directly in a step should only create a config node."""
monkeypatch.chdir(tmp_path)
config_data = {"setting": "direct", "value": 99}
config_path = _write_config(tmp_path, "direct_config.json", config_data)
class WorkflowWithDirectConfig(Workflow):
@step
async def step_with_direct_config(
self,
ev: StartEvent,
config: Annotated[
ConfigData, ResourceConfig(config_file="direct_config.json")
],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithDirectConfig()
graph = get_workflow_representation(workflow=wf)
assert len(_resource_nodes(graph)) == 0
resource_config_nodes = _resource_config_nodes(graph)
assert len(resource_config_nodes) == 1
config_node = resource_config_nodes[0]
assert config_node.type_name == "ConfigData"
assert config_node.config_file == config_path
assert config_node.config_value == config_data
edges = _find_edges(
graph,
source="step_with_direct_config",
target_prefix="resource_config_",
label="config",
)
assert len(edges) == 1
def test_resource_config_with_path_selector(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""ResourceConfig path selector is preserved in graph nodes."""
monkeypatch.chdir(tmp_path)
config_data = {"database": {"setting": "test", "value": 42}}
config_path = _write_config(tmp_path, "config.json", config_data)
def get_configured_client(
config: Annotated[
ConfigData,
ResourceConfig(config_file="config.json", path_selector="database"),
],
) -> DatabaseClient:
return DatabaseClient()
class WorkflowWithResourceConfig(Workflow):
@step
async def step_with_config(
self,
ev: StartEvent,
client: Annotated[DatabaseClient, Resource(get_configured_client)],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithResourceConfig()
graph = get_workflow_representation(workflow=wf)
resource_config_nodes = _resource_config_nodes(graph)
assert len(resource_config_nodes) == 1
config_node = resource_config_nodes[0]
assert config_node.config_file == config_path
assert config_node.path_selector == "database"
def test_resource_config_nodes_are_deduplicated(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Same resource config used by multiple resources appears once."""
monkeypatch.chdir(tmp_path)
_write_config(tmp_path, "config.json", {"setting": "test", "value": 42})
def get_client_one(
config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> DatabaseClient:
return DatabaseClient()
def get_client_two(
config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> DatabaseClient:
return DatabaseClient()
class WorkflowWithSharedConfig(Workflow):
@step
async def step_one(
self,
ev: StartEvent,
client: Annotated[DatabaseClient, Resource(get_client_one)],
) -> MiddleEvent:
return MiddleEvent()
@step
async def step_two(
self,
ev: MiddleEvent,
client: Annotated[DatabaseClient, Resource(get_client_two)],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithSharedConfig()
graph = get_workflow_representation(workflow=wf)
assert len(_resource_config_nodes(graph)) == 1
assert len(_resource_nodes(graph)) == 2
assert len(_find_edges(graph, target_prefix="resource_config_")) == 2
def test_multiple_different_resource_configs(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Multiple configs create distinct config nodes."""
monkeypatch.chdir(tmp_path)
db_path = _write_config(tmp_path, "db_config.json", {"setting": "db", "value": 1})
cache_path = _write_config(
tmp_path, "cache_config.json", {"setting": "cache", "value": 2}
)
def get_db_client(
config: Annotated[ConfigData, ResourceConfig(config_file="db_config.json")],
) -> DatabaseClient:
return DatabaseClient()
class CacheClient:
pass
def get_cache_client(
config: Annotated[ConfigData, ResourceConfig(config_file="cache_config.json")],
) -> CacheClient:
return CacheClient()
class WorkflowWithMultipleConfigs(Workflow):
@step
async def step_with_both(
self,
ev: StartEvent,
db: Annotated[DatabaseClient, Resource(get_db_client)],
cache: Annotated[CacheClient, Resource(get_cache_client)],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithMultipleConfigs()
graph = get_workflow_representation(workflow=wf)
config_files = {node.config_file for node in _resource_config_nodes(graph)}
assert config_files == {db_path, cache_path}
def test_filter_by_node_type_with_resource_config() -> None:
"""Test that filter_by_node_type works with resource_config nodes."""
graph = WorkflowGraph(
name="TestWorkflow",
nodes=[
WorkflowStepNode(id="step1", label="Step 1"),
WorkflowResourceNode(id="resource_123", label="Resource"),
WorkflowResourceConfigNode(
id="resource_config_456",
label="Config",
config_file="config.json",
),
],
edges=[
WorkflowGraphEdge(source="step1", target="resource_123", label="client"),
WorkflowGraphEdge(
source="resource_123", target="resource_config_456", label="config"
),
],
)
# Filter out resource_config nodes
filtered = graph.filter_by_node_type("resource_config")
assert len(filtered.nodes) == 2
node_types = {n.node_type for n in filtered.nodes}
assert node_types == {"step", "resource"}
# Edge from resource to config should be removed
# (no remaining node to connect to)
assert len(filtered.edges) == 1
assert filtered.edges[0].source == "step1"
assert filtered.edges[0].target == "resource_123"
def test_resource_config_label_and_description(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""ResourceConfig label and description are preserved in graph nodes."""
monkeypatch.chdir(tmp_path)
config_data = {"categories": ["invoice", "resume", "contract"]}
_write_config(tmp_path, "classify.json", config_data)
class WorkflowWithLabeledConfig(Workflow):
@step
async def classify_step(
self,
ev: StartEvent,
config: Annotated[
ConfigData,
ResourceConfig(
config_file="classify.json",
label="Document Classifier",
description="Configuration for document type classification",
),
],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithLabeledConfig()
graph = get_workflow_representation(workflow=wf)
config_nodes = _resource_config_nodes(graph)
assert len(config_nodes) == 1
config_node = config_nodes[0]
# Label should be used instead of type name
assert config_node.label == "Document Classifier"
assert config_node.description == "Configuration for document type classification"
# Type name should still be preserved
assert config_node.type_name == "ConfigData"
def test_resource_config_label_fallback(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""ResourceConfig without label falls back to type name."""
monkeypatch.chdir(tmp_path)
config_data = {"value": 123}
_write_config(tmp_path, "config.json", config_data)
class WorkflowWithUnlabeledConfig(Workflow):
@step
async def step(
self,
ev: StartEvent,
config: Annotated[ConfigData, ResourceConfig(config_file="config.json")],
) -> StopEvent:
return StopEvent(result="done")
wf = WorkflowWithUnlabeledConfig()
graph = get_workflow_representation(workflow=wf)
config_nodes = _resource_config_nodes(graph)
assert len(config_nodes) == 1
config_node = config_nodes[0]
# Label should fall back to type name when not specified
assert config_node.label == "ConfigData"
assert config_node.description is None