vellum-ai 1.11.2__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum_ee/workflows/display/utils/metadata.py

@@ -3,7 +3,7 @@
 import json
 import os
 from uuid import UUID
-from typing import Dict, Optional, Type
+from typing import Any, Dict, Optional, Tuple, Type
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.utils.files import virtual_open
@@ -33,30 +33,69 @@ def find_workflow_root_with_metadata(module_path: str) -> Optional[str]:
     return None
 
 
-def load_edges_to_id_mapping(module_path: str) -> Dict[str, str]:
+def _load_workflow_metadata(module_path: str) -> Tuple[Optional[str], Optional[Dict]]:
     """
-    Load edge path to ID mapping from metadata.json for a given module.
+    Load the full metadata.json content for a given module.
 
-    This function searches up the module hierarchy for metadata.json and extracts
-    the edges_to_id_mapping.
+    This function searches up the module hierarchy for metadata.json and loads its content.
 
     Args:
         module_path: The module path to search from (e.g., "workflows.my_workflow")
 
     Returns:
-        Dictionary mapping edge keys to their UUID strings
+        A tuple of (workflow_root, metadata_dict). Returns (None, None) if not found.
     """
     try:
-        root = find_workflow_root_with_metadata(module_path)
-        if not root:
-            return {}
-        file_path = os.path.join(root.replace(".", os.path.sep), "metadata.json")
+        workflow_root = find_workflow_root_with_metadata(module_path)
+        if not workflow_root:
+            return None, None
+        file_path = os.path.join(workflow_root.replace(".", os.path.sep), "metadata.json")
         with virtual_open(file_path) as f:
             data = json.load(f)
-            edges_map = data.get("edges_to_id_mapping")
-            return edges_map if isinstance(edges_map, dict) else {}
+            return workflow_root, data
     except Exception:
-        return {}
+        return None, None
+
+
+def _load_edges_mapping(module_path: str) -> Tuple[Optional[str], Dict[str, str]]:
+    """
+    Load edge path to ID mapping from metadata.json for a given module.
+
+    This function searches up the module hierarchy for metadata.json and extracts
+    the edges_to_id_mapping.
+
+    Args:
+        module_path: The module path to search from (e.g., "workflows.my_workflow")
+
+    Returns:
+        Dictionary mapping edge keys to their UUID strings
+    """
+    workflow_root, data = _load_workflow_metadata(module_path)
+    if data is None:
+        return None, {}
+    edges_map = data.get("edges_to_id_mapping")
+    if isinstance(edges_map, dict):
+        return workflow_root, edges_map
+    return workflow_root, {}
+
+
+def load_edges_to_id_mapping(module_path: str) -> Dict[str, str]:
+    _, mapping = _load_edges_mapping(module_path)
+    return mapping
+
+
+def _build_metadata_class_path(module: str, class_name: str, workflow_root: Optional[str]) -> str:
+    """
+    Format a class path so that it matches the metadata.json key structure.
+    """
+    if workflow_root:
+        if module == workflow_root:
+            return f".{class_name}"
+        prefix = f"{workflow_root}."
+        if module.startswith(prefix):
+            relative = module[len(prefix) :]
+            return f".{relative}.{class_name}" if relative else f".{class_name}"
+    return f"{module}.{class_name}"
 
 
 def get_trigger_edge_id(trigger_cls: Type, target_node: Type[BaseNode], module_path: str) -> Optional[str]:
@@ -71,9 +110,12 @@ def get_trigger_edge_id(trigger_cls: Type, target_node: Type[BaseNode], module_p
     Returns:
         The stable edge ID if found in metadata.json, None otherwise
     """
-    edges_mapping = load_edges_to_id_mapping(module_path)
-    trigger_path = f"{trigger_cls.__module__}.{trigger_cls.__name__}"
-    target_path = f"{target_node.__module__}.{target_node.__name__}.Trigger"
+    workflow_root, edges_mapping = _load_edges_mapping(module_path)
+    if not edges_mapping:
+        return None
+    trigger_path = _build_metadata_class_path(trigger_cls.__module__, trigger_cls.__name__, workflow_root)
+    target_base = _build_metadata_class_path(target_node.__module__, target_node.__name__, workflow_root)
+    target_path = f"{target_base}.Trigger"
     edge_key = f"{trigger_path}|{target_path}"
     return edges_mapping.get(edge_key)
 
@@ -89,9 +131,12 @@ def get_entrypoint_edge_id(target_node: Type[BaseNode], module_path: str) -> Opt
     Returns:
         The stable edge ID if found in metadata.json, None otherwise
     """
-    edges_mapping = load_edges_to_id_mapping(module_path)
+    workflow_root, edges_mapping = _load_edges_mapping(module_path)
+    if not edges_mapping:
+        return None
     manual_path = "vellum.workflows.triggers.manual.Manual"
-    target_path = f"{target_node.__module__}.{target_node.__name__}.Trigger"
+    target_base = _build_metadata_class_path(target_node.__module__, target_node.__name__, workflow_root)
+    target_path = f"{target_base}.Trigger"
     edge_key = f"{manual_path}|{target_path}"
     return edges_mapping.get(edge_key)
 
@@ -111,9 +156,13 @@ def get_regular_edge_id(
     Returns:
         The stable edge ID if found in metadata.json, None otherwise
     """
-    edges_mapping = load_edges_to_id_mapping(module_path)
-    source_path = f"{source_node_cls.__module__}.{source_node_cls.__name__}.Ports.{str(source_handle_id)}"
-    target_path = f"{target_node.__module__}.{target_node.__name__}.Trigger"
+    workflow_root, edges_mapping = _load_edges_mapping(module_path)
+    if not edges_mapping:
+        return None
+    source_base = _build_metadata_class_path(source_node_cls.__module__, source_node_cls.__name__, workflow_root)
+    source_path = f"{source_base}.Ports.{str(source_handle_id)}"
+    target_base = _build_metadata_class_path(target_node.__module__, target_node.__name__, workflow_root)
+    target_path = f"{target_base}.Trigger"
     edge_key = f"{source_path}|{target_path}"
     return edges_mapping.get(edge_key)
 
@@ -131,16 +180,32 @@ def load_dataset_row_index_to_id_mapping(module_path: str) -> Dict[int, str]:
     Returns:
         Dictionary mapping dataset row indices (as integers) to their ID strings
     """
-    try:
-        root = find_workflow_root_with_metadata(module_path)
-        if not root:
-            return {}
-        file_path = os.path.join(root.replace(".", os.path.sep), "metadata.json")
-        with virtual_open(file_path) as f:
-            data = json.load(f)
-            mapping = data.get("dataset_row_index_to_id_mapping")
-            if isinstance(mapping, dict):
-                return {int(k): v for k, v in mapping.items()}
-            return {}
-    except Exception:
+    _, data = _load_workflow_metadata(module_path)
+    if data is None:
+        return {}
+    mapping = data.get("dataset_row_index_to_id_mapping")
+    if isinstance(mapping, dict):
+        return {int(k): v for k, v in mapping.items()}
+    return {}
+
+
+def load_runner_config(module_path: str) -> Dict[str, Any]:
+    """
+    Load runner_config from metadata.json for a given module.
+
+    This function searches up the module hierarchy for metadata.json and extracts
+    the runner_config.
+
+    Args:
+        module_path: The module path to search from (e.g., "workflows.my_workflow")
+
+    Returns:
+        The runner_config dictionary if found in metadata.json, empty dict otherwise
+    """
+    _, data = _load_workflow_metadata(module_path)
+    if data is None:
         return {}
+    runner_config = data.get("runner_config")
+    if isinstance(runner_config, dict):
+        return runner_config
+    return {}

vellum_ee/workflows/display/utils/tests/test_metadata.py

@@ -0,0 +1,31 @@
+import pytest
+
+from vellum_ee.workflows.display.utils.metadata import _build_metadata_class_path
+
+
+def build_metadata_class_path_test_cases():
+    return [
+        # (module, class_name, workflow_root, expected)
+        ("my_workflow.nodes.my_node", "MyNode", "my_workflow", ".nodes.my_node.MyNode"),
+        ("my_workflow", "MyWorkflow", "my_workflow", ".MyWorkflow"),
+        ("my_workflow.triggers.scheduled", "ScheduleTrigger", "my_workflow", ".triggers.scheduled.ScheduleTrigger"),
+        ("external.module", "ExternalClass", "my_workflow", "external.module.ExternalClass"),
+        ("vellum.workflows.triggers.manual", "Manual", "my_workflow", "vellum.workflows.triggers.manual.Manual"),
+        ("my_workflow.nodes.my_node", "MyNode", None, "my_workflow.nodes.my_node.MyNode"),
+    ]
+
+
+@pytest.mark.parametrize(
+    "module, class_name, workflow_root, expected",
+    build_metadata_class_path_test_cases(),
+)
+def test_build_metadata_class_path(module, class_name, workflow_root, expected):
+    """
+    Tests that _build_metadata_class_path correctly formats class paths relative to workflow root.
+    """
+    # GIVEN a module path, class name, and workflow root
+    # WHEN we build the metadata class path
+    result = _build_metadata_class_path(module, class_name, workflow_root)
+
+    # THEN we should get the expected formatted path
+    assert result == expected

vellum_ee/workflows/display/utils/triggers.py

@@ -0,0 +1,153 @@
+from typing import Any, Dict, List, Type, cast
+
+from vellum.workflows.triggers.base import BaseTrigger
+from vellum.workflows.triggers.chat_message import ChatMessageTrigger
+from vellum.workflows.triggers.integration import IntegrationTrigger
+from vellum.workflows.triggers.manual import ManualTrigger
+from vellum.workflows.triggers.schedule import ScheduleTrigger
+from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+from vellum_ee.workflows.display.base import WorkflowTriggerType
+from vellum_ee.workflows.display.utils.vellum import compile_descriptor_annotation
+
+
+def get_trigger_type(trigger_class: Type[BaseTrigger]) -> WorkflowTriggerType:
+    """Get the WorkflowTriggerType for a trigger class."""
+    if issubclass(trigger_class, ManualTrigger):
+        return WorkflowTriggerType.MANUAL
+    elif issubclass(trigger_class, IntegrationTrigger):
+        return WorkflowTriggerType.INTEGRATION
+    elif issubclass(trigger_class, ScheduleTrigger):
+        return WorkflowTriggerType.SCHEDULED
+    elif issubclass(trigger_class, ChatMessageTrigger):
+        return WorkflowTriggerType.CHAT_MESSAGE
+    else:
+        raise ValueError(
+            f"Unknown trigger type: {trigger_class.__name__}. "
+            f"Please add it to the trigger type mapping in get_trigger_type_mapping()."
+        )
+
+
+def serialize_trigger_attributes(trigger_class: Type[BaseTrigger]) -> JsonArray:
+    """Serialize trigger attributes from attribute_references as VellumVariables."""
+    attribute_references = trigger_class.attribute_references().values()
+
+    def get_attribute_type(reference: Any) -> str:
+        # We can remove this type ignore with some mypy plugin changes
+        message_name = ChatMessageTrigger.message.name  # type: ignore[union-attr]
+        # For ChatMessageTrigger.message, always return ARRAY to maintain backwards compatibility
+        if issubclass(trigger_class, ChatMessageTrigger) and reference.name == message_name:
+            return "ARRAY"
+        try:
+            return primitive_type_to_vellum_variable_type(reference)
+        except ValueError:
+            return "JSON"
+
+    trigger_attributes: JsonArray = cast(
+        JsonArray,
+        [
+            cast(
+                JsonObject,
+                {
+                    "id": str(reference.id),
+                    "key": reference.name,
+                    "type": get_attribute_type(reference),
+                    "required": type(None) not in reference.types,
+                    "default": {
+                        "type": get_attribute_type(reference),
+                        "value": None,
+                    },
+                    "extensions": None,
+                    "schema": compile_descriptor_annotation(reference),
+                },
+            )
+            for reference in sorted(attribute_references, key=lambda ref: ref.name)
+        ],
+    )
+
+    return trigger_attributes
+
+
+def serialize_trigger_display_data(trigger_class: Type[BaseTrigger], trigger_type: WorkflowTriggerType) -> JsonObject:
+    """Serialize display_data from trigger's Display class."""
+    display_class = trigger_class.Display
+    display_data: JsonObject = {}
+
+    if hasattr(display_class, "label") and display_class.label is not None:
+        display_data["label"] = display_class.label
+
+    if (
+        hasattr(display_class, "x")
+        and display_class.x is not None
+        and hasattr(display_class, "y")
+        and display_class.y is not None
+    ):
+        display_data["position"] = {
+            "x": display_class.x,
+            "y": display_class.y,
+        }
+
+    if hasattr(display_class, "z_index") and display_class.z_index is not None:
+        display_data["z_index"] = display_class.z_index
+
+    if hasattr(display_class, "icon") and display_class.icon is not None:
+        display_data["icon"] = display_class.icon
+
+    if hasattr(display_class, "color") and display_class.color is not None:
+        display_data["color"] = display_class.color
+
+    if hasattr(display_class, "comment") and display_class.comment is not None:
+        display_data["comment"] = {
+            "value": display_class.comment.value,
+            "expanded": display_class.comment.expanded,
+        }
+
+    return display_data
+
+
+def serialize_trigger_definition(trigger_class: Type[BaseTrigger]) -> Dict[str, Any]:
+    """
+    Serialize a trigger class definition for use in node-definitions.json.
+
+    This produces a simplified trigger definition that includes:
+    - type: The WorkflowTriggerType enum value
+    - name: The trigger class name
+    - module: The module path as an array
+    - attributes: The trigger's attribute references
+    - display_data: Optional display metadata (label, icon, color)
+    """
+    trigger_type = get_trigger_type(trigger_class)
+
+    definition: Dict[str, Any] = {
+        "type": trigger_type.value,
+        "name": trigger_class.__name__,
+        "module": trigger_class.__module__.split("."),
+        "attributes": serialize_trigger_attributes(trigger_class),
+    }
+
+    display_data = serialize_trigger_display_data(trigger_class, trigger_type)
+
+    # Don't include display_data for manual triggers (consistent with workflow serialization)
+    if display_data and trigger_type != WorkflowTriggerType.MANUAL:
+        definition["display_data"] = display_data
+
+    return definition
+
+
+def get_all_trigger_classes() -> List[Type[BaseTrigger]]:
+    """
+    Get all trigger classes dynamically from the triggers module.
+
+    This mirrors the approach used by get_all_displayable_node_classes for nodes.
+    """
+    import vellum.workflows.triggers as triggers_module
+
+    trigger_classes = []
+    for class_name in triggers_module.__all__:
+        trigger_class = getattr(triggers_module, class_name)
+        # Skip BaseTrigger itself - we only want concrete trigger types
+        if trigger_class is BaseTrigger:
+            continue
+        if isinstance(trigger_class, type) and issubclass(trigger_class, BaseTrigger):
+            trigger_classes.append(trigger_class)
+    return trigger_classes

vellum_ee/workflows/display/utils/vellum.py

@@ -1,3 +1,4 @@
+import logging
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 from vellum.client.core.api_error import ApiError
@@ -14,14 +15,23 @@ from vellum.workflows.references.environment_variable import EnvironmentVariable
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.references.node import NodeReference
+from vellum.workflows.references.state_value import StateValueReference
+from vellum.workflows.references.trigger import TriggerAttributeReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.utils.functions import compile_annotation
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.utils.exceptions import InvalidInputReferenceError, UnsupportedSerializationException
+from vellum_ee.workflows.display.utils.exceptions import (
+    InvalidInputReferenceError,
+    InvalidOutputReferenceError,
+    UnsupportedSerializationException,
+)
 from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
+logger = logging.getLogger(__name__)
+
 
 class ConstantValuePointer(UniversalBaseModel):
     type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
@@ -78,6 +88,25 @@ class EnvironmentVariablePointer(UniversalBaseModel):
     data: EnvironmentVariableData
 
 
+class WorkflowStateData(UniversalBaseModel):
+    state_variable_id: str
+
+
+class WorkflowStatePointer(UniversalBaseModel):
+    type: Literal["WORKFLOW_STATE"] = "WORKFLOW_STATE"
+    data: WorkflowStateData
+
+
+class TriggerAttributeData(UniversalBaseModel):
+    trigger_id: str
+    attribute_id: str
+
+
+class TriggerAttributePointer(UniversalBaseModel):
+    type: Literal["TRIGGER_ATTRIBUTE"] = "TRIGGER_ATTRIBUTE"
+    data: TriggerAttributeData
+
+
 NodeInputValuePointerRule = Union[
     NodeOutputPointer,
     InputVariablePointer,
@@ -85,6 +114,8 @@ NodeInputValuePointerRule = Union[
     WorkspaceSecretPointer,
     ExecutionCounterPointer,
     EnvironmentVariablePointer,
+    WorkflowStatePointer,
+    TriggerAttributePointer,
 ]
 
 
@@ -108,6 +139,13 @@ def infer_vellum_variable_type(value: Any) -> VellumVariableType:
     return inferred_type
 
 
+def compile_descriptor_annotation(descriptor: BaseDescriptor) -> dict:
+    """Compile a BaseDescriptor's type annotation to a JSON schema."""
+    if not descriptor.types:
+        return {}
+    return compile_annotation(descriptor.normalized_type, {})
+
+
 def create_node_input_value_pointer_rule(
     value: Any, display_context: "WorkflowDisplayContext"
 ) -> NodeInputValuePointerRule:
@@ -132,7 +170,12 @@ def create_node_input_value_pointer_rule(
             data=NodeOutputData(node_id=str(upstream_node_display.node_id), output_id=str(output_display.id)),
         )
     if isinstance(value, LazyReference):
-        child_descriptor = get_child_descriptor(value, display_context)
+        try:
+            child_descriptor = get_child_descriptor(value, display_context)
+        except InvalidOutputReferenceError as e:
+            logger.warning("Failed to parse lazy reference '%s', skipping serialization", value.name)
+            display_context.add_validation_error(e)
+            raise UnsupportedSerializationException(f"Failed to parse lazy reference: {value.name}")
         return create_node_input_value_pointer_rule(child_descriptor, display_context)
     if isinstance(value, WorkflowInputReference):
         if value not in display_context.global_workflow_input_displays:
@@ -164,14 +207,25 @@ def create_node_input_value_pointer_rule(
             data=ExecutionCounterData(node_id=str(node_class_display.node_id)),
         )
     if isinstance(value, EnvironmentVariableReference):
-        if value.serialize_as_constant:
-            vellum_value = primitive_to_vellum_value(value.name)
-            return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
         return EnvironmentVariablePointer(
             data=EnvironmentVariableData(
                 environment_variable=value.name,
             ),
         )
+    if isinstance(value, StateValueReference):
+        state_value_display = display_context.global_state_value_displays[value]
+        return WorkflowStatePointer(
+            data=WorkflowStateData(
+                state_variable_id=str(state_value_display.id),
+            ),
+        )
+    if isinstance(value, TriggerAttributeReference):
+        return TriggerAttributePointer(
+            data=TriggerAttributeData(
+                trigger_id=str(value.trigger_class.__id__),
+                attribute_id=str(value.id),
+            ),
+        )
 
     if not isinstance(value, BaseDescriptor):
         vellum_value = primitive_to_vellum_value(value)