[KEP-2] Support user-injectable node preview functions

[ravi-kumar-pilla]

Context

• Related Issues/PRs: Issue #5190
• Authors: @ravi-kumar-pilla
• KEP shepherds: @ravi-kumar-pilla, @ElenaKhaustova

What are we trying to do?

Enable Kedro nodes to expose user-defined preview functions that allow visualization of complex, non-dataset objects directly in Kedro-Viz or other plugins.

The primary goal is to create an experimental feature that provides extensible preview mechanism for a kedro node while maintaining backward compatibility.

What this is NOT

• This is NOT about replacing existing dataset preview functionality
• This is NOT about modifying the core node execution logic
• This is NOT about creating a new visualization framework from scratch

Current Limitations

Currently, Kedro plugins (e.g., Kedro-Viz) can only preview standard dataset formats through the existing dataset preview system. Complex objects like:

• Compiled machine learning pipelines
• Graph structures (LangGraph, NetworkX graphs)
• Custom business objects
• Multi-modal data structures

...cannot be visualized directly.

Proposed Approach

Core Design

Based on the requirements from Issue #5190, the implementation will:

1. Start Simple: A preview function parameter in the existing Node definition
2. Experimental Feature: Implemented as an experimental feature to allow iteration. (Emit KedroExperimentalWarning inside node() when the user passes preview_fn)

Technical Architecture

This feature will:

1. Introduce an optional preview_fn parameter to Node definition that takes a callable which returns PreviewPayload contract.
2. Expose a node.preview() API for programmatic access
3. Support visual inspection of agentic workflows and ML pipeline intermediate objects

More info can be found in Appendix A

Key Features

1. Multiple Render Types: Support for Mermaid diagrams and existing dataset preview types (initial focus), with extensibility for JSON, images, Plotly charts, tables, and custom formats
2. Kedro-Viz Integration: Seamless display of previews in the pipeline visualization
3. Programmatic Access: Expose node.preview() API for external tools and debugging
4. Experimental Status: Clear warning to signal user that this is an experimental feature and bound to evolve.
5. PreviewPayload Contract: Standardized payload format for consistent rendering

Potential Drawbacks

• Increased API surface area (mitigated by experimental status)
• Initial limitation: no input parameters for preview functions (to be addressed in future iterations)

Who cares?

ML Engineers & Data Scientists:

• Direct visualization of complex ML pipeline structures and intermediate objects
• Better debugging and understanding of pipeline flows
• Enhanced collaboration through visual representations
• Support for agentic workflows and compiled models

Graph/Network Analysts:

• Native support for graph visualization in Kedro-Viz (LangGraph primary use case)
• Seamless integration of NetworkX, LangGraph, and other graph libraries
• Mermaid diagram support for flow visualization

Business Users:

• Improved pipeline comprehension through custom business object previews
• Better stakeholder communication with visual representations

DevOps/Platform Teams:

• Enhanced monitoring and debugging capabilities

Expected Improvements

• Developer Experience: Reduced context switching between Kedro-Viz and external visualization tools
• Pipeline Understanding: Visual representation of complex transformations
• Debugging Efficiency: Immediate visual feedback on pipeline components

Risks

1. Payload Size and timeout: Large preview payloads could impact Kedro-Viz performance. Either limit the size on Kedro-Viz or document any restrictions on payload size and timeout.
2. Maintenance Burden: Additional API surface to maintain and test

Timeline

~ 1.5 weeks

Appendix A: Proposed API Changes

Preview Function Parameter:

# kedro/pipeline/preview_types.py
from typing import Any, Literal, Callable
import inspect

# This will establish a contract, so anyone implementing a node preview function should return this object
# Also makes plugin developers aware of what to expect from node preview function
class PreviewPayload():
    kind: Literal["mermaid", "json", "text", "image", "plotly", "table", "custom"]
    content: str
    meta: dict[str, Any] | None

# Node class with preview_fn
class Node:
    def __init__(self, func, inputs, outputs, name=None, tags=None, preview_fn: Optional[Callable[..., PreviewPayload]] = None):
        self.func = func
        self.inputs = inputs
        self.outputs = outputs
        self.name = name
        self.tags = tags or ()
        self.preview_fn = preview_fn

    def preview(self, *args, **kwargs) -> Optional[PreviewPayload]:
        if not self.preview_fn:
            return None
        return self.preview_fn(*args, **kwargs)

# Usage
def preview_graph() -> PreviewPayload:
    return PreviewPayload(kind="mermaid", content="graph TD; A-->B")
   
node(
    func=compile_graph,
    outputs="compiled_graph",
    preview_fn=preview_graph  # Simple function, no args support
)

Appendix B. Optional Design Sketch

Appendix C: Optional Rejected Designs

NodePreview Instance:

1. Extend kedro.pipeline.node.Node

# kedro/pipeline/node.py
from typing import Any, Callable
import inspect
from kedro.pipeline.preview_types import NodePreview, PreviewPayload

class Node:
    def __init__(self, func, inputs, outputs, name=None, tags=None, preview: NodePreview | None = None):
        self.func = func
        self.inputs = inputs
        self.outputs = outputs
        self.name = name
        self.tags = tags or ()
        self._preview = preview

    @property
    def has_preview(self) -> bool:
        """Check if node has preview capability."""
        return self._preview is not None

    def preview(self, *args: Any, **kwargs: Any) -> PreviewPayload | None:
        """Execute preview with intelligent argument merging."""
        if not self._preview:
            return None

        # Merge args and kwargs from NodePreview and method call
        merged_args = self._preview.args + args
        merged_kwargs = {**self._preview.kwargs, **kwargs}

        # Use signature inspection to call the function intelligently
        sig = inspect.signature(self._preview.fn)
        try:
            bound = sig.bind(*merged_args, **merged_kwargs)
            bound.apply_defaults()
            return self._preview.fn(*bound.args, **bound.kwargs)
        except TypeError:
            # Fallback to direct call if signature binding fails
            return self._preview.fn(*merged_args, **merged_kwargs)

2. Node constructor helper

# kedro/pipeline/__init__.py
from kedro.pipeline.node import Node
from kedro.pipeline.preview_types import NodePreview

def node(func, inputs, outputs, *, name=None, tags=None, preview=None):
    return Node(func, inputs, outputs, name, tags, preview)

3. Example usage (LangGraph)

# project/pipelines/agentic/nodes.py
from kedro.pipeline.preview_types import PreviewPayload, NodePreview

def compile_langgraph(cfg):
    # user code that returns compiled graph object with draw_mermaid()
    graph = build_langgraph(cfg)
    return graph

def visualize_langgraph(theme="default", format="mermaid"):
    """Preview function with intelligent args/kwargs merging."""
    content = "graph TD\n  A[Config] --> B[LangGraph]\n  B --> C[Compiled]"
    if theme != "default":
        content = f"%%{{init: {{'theme': '{theme}'}}}}%%\n{content}"
    return {"kind": "mermaid", "content": content, "meta": {"theme": theme}}

# pipeline definition
from kedro.pipeline import node, Pipeline

pipeline = Pipeline([
    node(
        func=compile_langgraph,
        inputs="config",
        outputs="compiled_graph",
        name="compile_langgraph",
        preview=NodePreview(
            fn=visualize_langgraph,
            kwargs={"theme": "dark", "format": "mermaid"}
        ),
    )
])

# Usage: Both NodePreview kwargs and runtime preview() kwargs are merged
# node.preview()  # Uses theme="dark" from NodePreview
# node.preview(theme="light")  # Overrides with theme="light"

4. Usage Patterns

Simple examples showing args/kwargs support:

# Pattern 1: Simple preview function
def simple_preview():
    return {"kind": "text", "content": "This node generates a graph"}

# Pattern 2: Preview with custom arguments
def themed_preview(theme="light", show_meta=True):
    content = f"Generated visualization with {theme} theme"
    result = {"kind": "text", "content": content}
    if show_meta:
        result["meta"] = {"theme": theme, "timestamp": "2024-01-01"}
    return result

# Usage examples
node(
    func=my_func,
    outputs="result",
    preview=NodePreview(fn=simple_preview)
)

node(
    func=my_func,
    outputs="result",
    preview=NodePreview(
        fn=themed_preview,
        kwargs={"theme": "dark", "show_meta": False}
    )
)

# Updated node function signature
def node(
    func: Callable,
    inputs: str | list[str] | dict[str, str] | None = None,
    outputs: str | list[str] | dict[str, str] | None = None, 
    name: str | None = None,
    tags: str | Iterable[str] | None = None,
    confirms: str | list[str] | None = None,
    namespace: str | None = None,
    preview_fn: Callable[[], PreviewPayload] | None = None,  # NEW - Experimental
) -> Node:
    ...

Decorator-Based:

# Decorator to attach preview to function
def with_node_preview(preview_fn: Callable[..., PreviewPayload]):
    """Decorator to attach a preview function to a node function.

    Args:
        preview_fn: A function that returns a PreviewPayload instance.

    Example:
        # Preview function with no arguments
        def preview_graph() -> PreviewPayload:
            return PreviewPayload(kind="mermaid", content="graph TD; A-->B")

        # Preview function with arguments
        def preview_with_args(data) -> PreviewPayload:
            return PreviewPayload(kind="text", content=f"Preview: {data}")

        @with_node_preview(preview_graph)
        def compile_graph(config):
            return build_graph(config)
    """

    def decorator(func):
        func.__kedro_node_preview_fn__ = preview_fn
        return func

    return decorator

def preview(self, *args: Any, **kwargs: Any) -> PreviewPayload | None:
      """Generate a preview of the node.

      Args:
          *args: Positional arguments to pass to the preview function.
          **kwargs: Keyword arguments to pass to the preview function.

      Returns:
          Preview payload if preview is configured, None otherwise.
      """
      preview_fn = getattr(self.func, "__kedro_node_preview_fn__", None)
      if not preview_fn:
          return None

      # Use signature inspection to call the function intelligently
      sig = inspect.signature(preview_fn)
      try:
          # Try to bind the arguments to the function signature
          bound = sig.bind(*args, **kwargs)
          bound.apply_defaults()
          return preview_fn(*bound.args, **bound.kwargs)
      except TypeError:
          # If binding fails (e.g., function takes no args but args were provided),
          # try calling with no arguments
          try:
              return preview_fn()
          except TypeError:
              # If that also fails, try calling with all provided arguments
              return preview_fn(*args, **kwargs)
                
# Usage
def visualize_companies_graph() -> PreviewPayload:
    """
    Visualization function that returns a Mermaid diagram of the intent detection workflow.
    
    Args:
        result: The result from detect_intent function
    
    Returns:
        Mermaid diagram representation of the intent detection flow.
    """
    return PreviewPayload(kind = "mermaid", content = """graph TD
A[START] --> B[detect_intent]
B --> C{Intent Classification}
C -->|clarification_needed| D[clarify_intent]
C -->|general_question| E[update_context]
C -->|claim_new| E
C -->|existing_claim_question| E
D --> E
E --> F[END]
style A fill:#e1f5fe
style F fill:#f3e5f5
style C fill:#fff3e0
style D fill:#ffebee
style E fill:#e8f5e8"""
    )

@with_node_preview(visualize_companies_graph)
def apply_types_to_companies(companies: pd.DataFrame) -> pd.DataFrame:
    """Preprocesses the data for companies.

    Args:
        companies: Raw data.
    Returns:
        Preprocessed data, with `company_rating` converted to a float and
        `iata_approved` converted to boolean.
    """
    companies["iata_approved"] = _is_true(companies["iata_approved"])
    companies["company_rating"] = (
        companies["company_rating"].str.replace("%", "").astype(float) / 100
    )
    return companies

[ravi-kumar-pilla]

This proposal will be implemented in the next sprint. Please follow - #5317 for the implementation.