run-llama/workflows-py
1
0
Fork 0
mirror of https://github.com/run-llama/workflows-py.git synced 2026-08-24 20:01:34 -04:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

refact: split out workflow representations, to prep for separate server package (#274)

Browse Source
This commit is contained in:
Adrian Lyjak
2026-01-13 13:26:54 -05:00
committed by GitHub
parent 7a85c96d68
commit 0d72b4ded0
8 changed files with 322 additions and 264 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/some-zebras-wash.md
+6
View File
@@ -0,0 +1,6 @@
---
"llama-index-utils-workflow": minor
"llama-index-workflows": minor
---
reorganize workflow graph representation types
packages/llama-index-utils-workflow/src/llama_index/utils/workflow/__init__.py
+5 -5
View File
@@ -18,15 +18,15 @@ from workflows.events import (
StopEvent,
)
from workflows.handler import WorkflowHandler
from workflows.protocol import (
from workflows.representation import (
WorkflowGenericNode,
WorkflowGraph,
WorkflowGraphEdge,
WorkflowGraphNode,
WorkflowResourceNode,
)
from workflows.representation_utils import (
extract_workflow_structure as _extract_workflow_structure,
from workflows.representation import (
get_workflow_representation as _get_workflow_representation,
)
from workflows.runtime.types.results import AddCollectedEvent, StepWorkerResult
from workflows.runtime.types.ticks import TickAddEvent, TickStepResult, WorkflowTick
@@ -511,7 +511,7 @@ def draw_all_possible_flows(
max_label_length: Maximum label length before truncation (None = no limit)
"""
graph = _extract_workflow_structure(workflow)
graph = _get_workflow_representation(workflow)
_render_pyvis(graph, filename, notebook, max_label_length)
@@ -532,7 +532,7 @@ def draw_all_possible_flows_mermaid(
The Mermaid diagram as a string
"""
graph = _extract_workflow_structure(workflow)
graph = _get_workflow_representation(workflow)
return _render_mermaid(graph, filename, max_label_length)
packages/llama-index-workflows/src/workflows/protocol/__init__.py
+3 -244
View File
@@ -1,10 +1,11 @@
from __future__ import annotations
from typing import Any, Literal, Union
from typing import Any, Literal
from pydantic import BaseModel, Field
from workflows.protocol.serializable_events import EventEnvelopeWithMetadata
from workflows.representation import WorkflowGraph
# Shared protocol types between client and server
@@ -68,242 +69,9 @@ class WorkflowGraphResponse(BaseModel):
graph: WorkflowGraph
class WorkflowNodeBase(BaseModel):
"""Base class for all workflow graph nodes."""
id: str = Field(description="Unique identifier for the node")
label: str = Field(description="Display text for the node")
def truncated_label(self, max_length: int) -> str:
"""Get truncated label for visualization (adds * suffix if truncated)."""
if len(self.label) <= max_length:
return self.label
return f"{self.label[: max_length - 1]}*"
class WorkflowStepNode(WorkflowNodeBase):
"""A workflow step node representing a function decorated with @step."""
node_type: Literal["step"] = Field(
default="step", description="Discriminator field for node type"
)
description: str | None = Field(
default=None,
description="Documentation string extracted from the step function",
)
class WorkflowEventNode(WorkflowNodeBase):
"""An event node representing an Event class that flows between steps."""
node_type: Literal["event"] = Field(
default="event", description="Discriminator field for node type"
)
event_type: str = Field(
description="The event class name (e.g., 'StartEvent', 'MyCustomEvent')"
)
event_types: list[str] = Field(
description="Event class inheritance chain for subclass checking. "
"First element is the class itself, followed by parent Event subclasses."
)
event_schema: dict[str, Any] | None = Field(
default=None,
description="Pydantic JSON schema for the event type",
)
def is_subclass_of(self, *type_names: str) -> bool:
"""Check if this node's event_type is a subclass of any of the given types."""
return any(name in self.event_types for name in type_names)
class WorkflowExternalNode(WorkflowNodeBase):
"""An external node representing human-in-the-loop or external system interaction."""
node_type: Literal["external"] = Field(
default="external", description="Discriminator field for node type"
)
class WorkflowResourceNode(WorkflowNodeBase):
"""A resource node representing an injected dependency (e.g., database client, API client)."""
node_type: Literal["resource"] = Field(
default="resource", description="Discriminator field for node type"
)
type_name: str | None = Field(
default=None,
description="The type annotation of the resource (e.g., 'DatabaseClient', 'AsyncLlamaCloud')",
)
getter_name: str | None = Field(
default=None,
description="Name of the factory function that creates the resource",
)
source_file: str | None = Field(
default=None,
description="Absolute path to the source file containing the getter function",
)
source_line: int | None = Field(
default=None, description="Line number where the getter function is defined"
)
description: str | None = Field(
default=None,
description="Documentation string extracted from the getter function",
)
class WorkflowGenericNode(WorkflowNodeBase):
"""A generic node for custom visualization types not covered by standard node types.
Used for agent visualization (node_type='agent', 'tool', 'workflow_agent', etc.)
and other custom extensions. Supports optional event_type fields for type checking.
"""
node_type: str = Field(
description="Custom node type string (e.g., 'agent', 'tool', 'workflow_base')"
)
event_type: str | None = Field(
default=None,
description="Optional type name for nodes that support inheritance checking (e.g., agent types)",
)
event_types: list[str] | None = Field(
default=None,
description="Optional inheritance chain for subclass checking, similar to WorkflowEventNode",
)
def is_subclass_of(self, *type_names: str) -> bool:
"""Check if this node's event_type is a subclass of any of the given types."""
if not self.event_types:
return False
return any(name in self.event_types for name in type_names)
# Union type for workflow graph nodes
# Pydantic will try to match against types in order; WorkflowGenericNode is last as catch-all
WorkflowGraphNode = Union[
WorkflowStepNode,
WorkflowEventNode,
WorkflowExternalNode,
WorkflowResourceNode,
WorkflowGenericNode,
]
class WorkflowGraphEdge(BaseModel):
"""A directed edge connecting two nodes in the workflow graph."""
source: str = Field(description="ID of the source node (where the edge originates)")
target: str = Field(description="ID of the target node (where the edge points to)")
label: str | None = Field(
default=None,
description="Optional edge label, used for resource edges to show the variable name",
)
class WorkflowGraph(BaseModel):
"""Complete workflow graph structure containing all nodes and edges."""
nodes: list[WorkflowGraphNode] = Field(
description="All nodes in the workflow graph"
)
edges: list[WorkflowGraphEdge] = Field(
description="All directed edges connecting the nodes"
)
description: str | None = Field(
default=None,
description="Documentation string extracted from the workflow class",
)
def filter_by_node_type(self, *node_types: str) -> WorkflowGraph:
"""Create a simplified graph by removing nodes of specified types.
Edges passing through filtered nodes are resolved:
Node1 -> FilteredNode -> Node2 becomes Node1 -> Node2
Args:
*node_types: One or more node type strings to filter out
(e.g., "event", "resource", "step", "external")
Returns:
A new WorkflowGraph with the specified node types removed
and edges resolved through them.
"""
filter_types = set(node_types)
# Identify nodes to filter out
filtered_node_ids: set[str] = set()
for node in self.nodes:
if node.node_type in filter_types:
filtered_node_ids.add(node.id)
# Keep remaining nodes
remaining_nodes = [n for n in self.nodes if n.id not in filtered_node_ids]
remaining_node_ids = {n.id for n in remaining_nodes}
# Build outgoing edge map and node lookup
outgoing_map: dict[str, list[WorkflowGraphEdge]] = {}
for edge in self.edges:
outgoing_map.setdefault(edge.source, []).append(edge)
node_by_id: dict[str, WorkflowGraphNode] = {n.id: n for n in self.nodes}
def resolve_targets(
from_id: str,
first_filtered_label: str | None,
visited: set[str],
) -> list[tuple[str, str | None]]:
"""Find remaining nodes reachable from from_id, through filtered nodes."""
results: list[tuple[str, str | None]] = []
for edge in outgoing_map.get(from_id, []):
target = edge.target
if target in visited:
continue
if target in remaining_node_ids:
# Use the first filtered node's label, or the edge label if direct
label = (
first_filtered_label
if first_filtered_label is not None
else edge.label
)
results.append((target, label))
elif target in filtered_node_ids:
# Follow through filtered node, capturing its label if first
visited.add(target)
filtered_node = node_by_id[target]
label = (
first_filtered_label
if first_filtered_label is not None
else filtered_node.label
)
results.extend(resolve_targets(target, label, visited))
return results
# Build new edges
new_edges: list[WorkflowGraphEdge] = []
seen_edges: set[tuple[str, str]] = set()
for source_id in remaining_node_ids:
for target_id, label in resolve_targets(source_id, None, set()):
edge_key = (source_id, target_id)
if edge_key not in seen_edges:
seen_edges.add(edge_key)
new_edges.append(
WorkflowGraphEdge(
source=source_id,
target=target_id,
label=label,
)
)
return WorkflowGraph(
nodes=remaining_nodes,
edges=new_edges,
description=self.description,
)
__all__ = [
"Status",
"is_status_completed",
"HandlerData",
"HandlersListResponse",
"HealthResponse",
@@ -313,13 +81,4 @@ __all__ = [
"WorkflowSchemaResponse",
"WorkflowEventsListResponse",
"WorkflowGraphResponse",
"WorkflowNodeBase",
"WorkflowStepNode",
"WorkflowEventNode",
"WorkflowExternalNode",
"WorkflowResourceNode",
"WorkflowGenericNode",
"WorkflowGraphNode",
"WorkflowGraphEdge",
"WorkflowGraph",
]
packages/llama-index-workflows/src/workflows/representation/__init__.py
+27
View File
@@ -0,0 +1,27 @@
from workflows.representation.build import get_workflow_representation
from workflows.representation.types import (
WorkflowEventNode,
WorkflowExternalNode,
WorkflowGenericNode,
WorkflowGraph,
WorkflowGraphEdge,
WorkflowGraphNode,
WorkflowNodeBase,
WorkflowResourceNode,
WorkflowStepNode,
)
__all__ = [
# Types
"WorkflowNodeBase",
"WorkflowStepNode",
"WorkflowEventNode",
"WorkflowExternalNode",
"WorkflowResourceNode",
"WorkflowGenericNode",
"WorkflowGraphNode",
"WorkflowGraphEdge",
"WorkflowGraph",
# Functions
"get_workflow_representation",
]
packages/llama-index-workflows/src/workflows/representation_utils.py → packages/llama-index-workflows/src/workflows/representation/build.py
+17 -3
View File
@@ -11,7 +11,7 @@ from workflows.events import (
InputRequiredEvent,
StopEvent,
)
from workflows.protocol import (
from workflows.representation.types import (
WorkflowEventNode,
WorkflowExternalNode,
WorkflowGraph,
@@ -92,8 +92,19 @@ def _create_resource_node(resource_def: ResourceDefinition) -> WorkflowResourceN
)
def extract_workflow_structure(workflow: Workflow) -> WorkflowGraph:
"""Extract workflow structure into a graph representation."""
def get_workflow_representation(workflow: Workflow) -> WorkflowGraph:
"""Build a graph representation of a workflow's structure.
Extracts the workflow's steps, events, and resources into a WorkflowGraph
that can be used for visualization or analysis.
Args:
workflow: The workflow instance to build a representation for.
Returns:
A WorkflowGraph containing nodes for steps, events, resources,
and external interactions, with edges showing the data flow.
"""
# Get workflow steps
steps: dict[str, StepFunction] = get_steps_from_class(workflow)
if not steps:
@@ -239,3 +250,6 @@ def extract_workflow_structure(workflow: Workflow) -> WorkflowGraph:
workflow_description = inspect.getdoc(workflow)
return WorkflowGraph(nodes=nodes, edges=edges, description=workflow_description)
__all__ = ["get_workflow_representation"]
packages/llama-index-workflows/src/workflows/representation/types.py
+252
View File
@@ -0,0 +1,252 @@
from __future__ import annotations
from typing import Any, Literal, Union
from pydantic import BaseModel, Field
class WorkflowNodeBase(BaseModel):
"""Base class for all workflow graph nodes."""
id: str = Field(description="Unique identifier for the node")
label: str = Field(description="Display text for the node")
def truncated_label(self, max_length: int) -> str:
"""Get truncated label for visualization (adds * suffix if truncated)."""
if len(self.label) <= max_length:
return self.label
return f"{self.label[: max_length - 1]}*"
class WorkflowStepNode(WorkflowNodeBase):
"""A workflow step node representing a function decorated with @step."""
node_type: Literal["step"] = Field(
default="step", description="Discriminator field for node type"
)
description: str | None = Field(
default=None,
description="Documentation string extracted from the step function",
)
class WorkflowEventNode(WorkflowNodeBase):
"""An event node representing an Event class that flows between steps."""
node_type: Literal["event"] = Field(
default="event", description="Discriminator field for node type"
)
event_type: str = Field(
description="The event class name (e.g., 'StartEvent', 'MyCustomEvent')"
)
event_types: list[str] = Field(
description="Event class inheritance chain for subclass checking. "
"First element is the class itself, followed by parent Event subclasses."
)
event_schema: dict[str, Any] | None = Field(
default=None,
description="Pydantic JSON schema for the event type",
)
def is_subclass_of(self, *type_names: str) -> bool:
"""Check if this node's event_type is a subclass of any of the given types."""
return any(name in self.event_types for name in type_names)
class WorkflowExternalNode(WorkflowNodeBase):
"""An external node representing human-in-the-loop or external system interaction."""
node_type: Literal["external"] = Field(
default="external", description="Discriminator field for node type"
)
class WorkflowResourceNode(WorkflowNodeBase):
"""A resource node representing an injected dependency (e.g., database client, API client)."""
node_type: Literal["resource"] = Field(
default="resource", description="Discriminator field for node type"
)
type_name: str | None = Field(
default=None,
description="The type annotation of the resource (e.g., 'DatabaseClient', 'AsyncLlamaCloud')",
)
getter_name: str | None = Field(
default=None,
description="Name of the factory function that creates the resource",
)
source_file: str | None = Field(
default=None,
description="Absolute path to the source file containing the getter function",
)
source_line: int | None = Field(
default=None, description="Line number where the getter function is defined"
)
description: str | None = Field(
default=None,
description="Documentation string extracted from the getter function",
)
class WorkflowGenericNode(WorkflowNodeBase):
"""A generic node for custom visualization types not covered by standard node types.
Used for agent visualization (node_type='agent', 'tool', 'workflow_agent', etc.)
and other custom extensions. Supports optional event_type fields for type checking.
"""
node_type: str = Field(
description="Custom node type string (e.g., 'agent', 'tool', 'workflow_base')"
)
event_type: str | None = Field(
default=None,
description="Optional type name for nodes that support inheritance checking (e.g., agent types)",
)
event_types: list[str] | None = Field(
default=None,
description="Optional inheritance chain for subclass checking, similar to WorkflowEventNode",
)
def is_subclass_of(self, *type_names: str) -> bool:
"""Check if this node's event_type is a subclass of any of the given types."""
if not self.event_types:
return False
return any(name in self.event_types for name in type_names)
# Union type for workflow graph nodes
# Pydantic will try to match against types in order; WorkflowGenericNode is last as catch-all
WorkflowGraphNode = Union[
WorkflowStepNode,
WorkflowEventNode,
WorkflowExternalNode,
WorkflowResourceNode,
WorkflowGenericNode,
]
class WorkflowGraphEdge(BaseModel):
"""A directed edge connecting two nodes in the workflow graph."""
source: str = Field(description="ID of the source node (where the edge originates)")
target: str = Field(description="ID of the target node (where the edge points to)")
label: str | None = Field(
default=None,
description="Optional edge label, used for resource edges to show the variable name",
)
class WorkflowGraph(BaseModel):
"""Complete workflow graph structure containing all nodes and edges."""
nodes: list[WorkflowGraphNode] = Field(
description="All nodes in the workflow graph"
)
edges: list[WorkflowGraphEdge] = Field(
description="All directed edges connecting the nodes"
)
description: str | None = Field(
default=None,
description="Documentation string extracted from the workflow class",
)
def filter_by_node_type(self, *node_types: str) -> WorkflowGraph:
"""Create a simplified graph by removing nodes of specified types.
Edges passing through filtered nodes are resolved:
Node1 -> FilteredNode -> Node2 becomes Node1 -> Node2
Args:
*node_types: One or more node type strings to filter out
(e.g., "event", "resource", "step", "external")
Returns:
A new WorkflowGraph with the specified node types removed
and edges resolved through them.
"""
filter_types = set(node_types)
# Identify nodes to filter out
filtered_node_ids: set[str] = set()
for node in self.nodes:
if node.node_type in filter_types:
filtered_node_ids.add(node.id)
# Keep remaining nodes
remaining_nodes = [n for n in self.nodes if n.id not in filtered_node_ids]
remaining_node_ids = {n.id for n in remaining_nodes}
# Build outgoing edge map and node lookup
outgoing_map: dict[str, list[WorkflowGraphEdge]] = {}
for edge in self.edges:
outgoing_map.setdefault(edge.source, []).append(edge)
node_by_id: dict[str, WorkflowGraphNode] = {n.id: n for n in self.nodes}
def resolve_targets(
from_id: str,
first_filtered_label: str | None,
visited: set[str],
) -> list[tuple[str, str | None]]:
"""Find remaining nodes reachable from from_id, through filtered nodes."""
results: list[tuple[str, str | None]] = []
for edge in outgoing_map.get(from_id, []):
target = edge.target
if target in visited:
continue
if target in remaining_node_ids:
# Use the first filtered node's label, or the edge label if direct
label = (
first_filtered_label
if first_filtered_label is not None
else edge.label
)
results.append((target, label))
elif target in filtered_node_ids:
# Follow through filtered node, capturing its label if first
visited.add(target)
filtered_node = node_by_id[target]
label = (
first_filtered_label
if first_filtered_label is not None
else filtered_node.label
)
results.extend(resolve_targets(target, label, visited))
return results
# Build new edges
new_edges: list[WorkflowGraphEdge] = []
seen_edges: set[tuple[str, str]] = set()
for source_id in remaining_node_ids:
for target_id, label in resolve_targets(source_id, None, set()):
edge_key = (source_id, target_id)
if edge_key not in seen_edges:
seen_edges.add(edge_key)
new_edges.append(
WorkflowGraphEdge(
source=source_id,
target=target_id,
label=label,
)
)
return WorkflowGraph(
nodes=remaining_nodes,
edges=new_edges,
description=self.description,
)
__all__ = [
"WorkflowNodeBase",
"WorkflowStepNode",
"WorkflowEventNode",
"WorkflowExternalNode",
"WorkflowResourceNode",
"WorkflowGenericNode",
"WorkflowGraphNode",
"WorkflowGraphEdge",
"WorkflowGraph",
]
packages/llama-index-workflows/src/workflows/server/server.py
+2 -2
View File
@@ -52,7 +52,7 @@ from workflows.protocol.serializable_events import (
EventEnvelopeWithMetadata,
EventValidationError,
)
from workflows.representation_utils import extract_workflow_structure
from workflows.representation import get_workflow_representation
from workflows.server.abstract_workflow_store import (
AbstractWorkflowStore,
HandlerQuery,
@@ -653,7 +653,7 @@ class WorkflowServer:
"""
workflow = self._extract_workflow(request)
try:
workflow_graph = extract_workflow_structure(workflow.workflow)
workflow_graph = get_workflow_representation(workflow.workflow)
except Exception as e:
raise HTTPException(
detail=f"Error while getting JSON workflow representation: {e}",
packages/llama-index-workflows/tests/test_representation_utils.py
+10 -10
View File
@@ -3,15 +3,15 @@ from typing import Annotated
import pytest
from workflows.decorators import step
from workflows.events import Event, StartEvent, StopEvent
from workflows.protocol import (
from workflows.representation import (
WorkflowEventNode,
WorkflowExternalNode,
WorkflowGraph,
WorkflowGraphEdge,
WorkflowResourceNode,
WorkflowStepNode,
get_workflow_representation,
)
from workflows.representation_utils import extract_workflow_structure
from workflows.resource import Resource
from workflows.workflow import Workflow
@@ -70,9 +70,9 @@ def ground_truth_repr() -> WorkflowGraph:
)
def test_extract_workflow_structure(ground_truth_repr: WorkflowGraph) -> None:
def test_get_workflow_representation(ground_truth_repr: WorkflowGraph) -> None:
wf = DummyWorkflow()
graph = extract_workflow_structure(workflow=wf)
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"]
@@ -174,10 +174,10 @@ class WorkflowWithResources(Workflow):
return StopEvent(result="done")
def test_extract_workflow_structure_with_resources() -> None:
def test_get_workflow_representation_with_resources() -> None:
"""Test that resource nodes are extracted from workflow with resources."""
wf = WorkflowWithResources()
graph = extract_workflow_structure(workflow=wf)
graph = get_workflow_representation(workflow=wf)
# Should have resource nodes
resource_nodes = [n for n in graph.nodes if isinstance(n, WorkflowResourceNode)]
@@ -196,7 +196,7 @@ def test_extract_workflow_structure_with_resources() -> 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 = extract_workflow_structure(workflow=wf)
graph = get_workflow_representation(workflow=wf)
# Find edges to resource nodes
resource_edges = [e for e in graph.edges if e.target.startswith("resource_")]
@@ -235,7 +235,7 @@ def test_resource_nodes_are_deduplicated() -> None:
return StopEvent(result="done")
wf = WorkflowWithSharedResource()
graph = extract_workflow_structure(workflow=wf)
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)]
@@ -270,7 +270,7 @@ def test_multiple_different_resources() -> None:
return StopEvent(result="done")
wf = WorkflowWithMultipleResources()
graph = extract_workflow_structure(workflow=wf)
graph = get_workflow_representation(workflow=wf)
# Should have two different resource nodes
resource_nodes = [n for n in graph.nodes if isinstance(n, WorkflowResourceNode)]
@@ -326,7 +326,7 @@ def test_resource_node_serialization() -> None:
def test_graph_with_resources() -> None:
"""Test that workflow graph with resources is correct."""
wf = WorkflowWithResources()
graph = extract_workflow_structure(workflow=wf)
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)]