mirror of
https://github.com/run-llama/workflows-py.git
synced 2026-07-21 12:15:24 -04:00
Add resource config node support to workflow representation (#294)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user