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

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
@@ -8,18 +9,29 @@ from vellum.client.types.vellum_variable_type import VellumVariableType
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
+from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.references import OutputReference, WorkflowInputReference
+from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 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 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"
@@ -67,12 +79,43 @@ class ExecutionCounterPointer(UniversalBaseModel):
     data: ExecutionCounterData
 
 
+class EnvironmentVariableData(UniversalBaseModel):
+    environment_variable: str
+
+
+class EnvironmentVariablePointer(UniversalBaseModel):
+    type: Literal["ENVIRONMENT_VARIABLE"] = "ENVIRONMENT_VARIABLE"
+    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,
     ConstantValuePointer,
     WorkspaceSecretPointer,
     ExecutionCounterPointer,
+    EnvironmentVariablePointer,
+    WorkflowStatePointer,
+    TriggerAttributePointer,
 ]
 
 
@@ -96,25 +139,51 @@ 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:
     if isinstance(value, OutputReference):
         if value not in display_context.global_node_output_displays:
             if issubclass(value.outputs_class, BaseNode.Outputs):
-                raise ValueError(f"Reference to node '{value.outputs_class._node_class.__name__}' not found in graph.")
+                raise ValueError(
+                    f"Reference to node '{value.outputs_class.__parent_class__.__name__}' not found in graph."
+                )
 
             raise ValueError(f"Reference to outputs '{value.outputs_class.__qualname__}' is invalid.")
 
-        upstream_node, output_display = display_context.global_node_output_displays[value]
-        upstream_node_display = display_context.global_node_displays[upstream_node]
+        output_display = display_context.global_node_output_displays[value]
+
+        upstream_node_class = value.outputs_class.__parent_class__
+        if not issubclass(upstream_node_class, BaseNode):
+            raise ValueError(f"Output references must be to a node, not {upstream_node_class}")
+        unadorned_upstream_node_class = get_unadorned_node(upstream_node_class)
+        upstream_node_display = display_context.global_node_displays[unadorned_upstream_node_class]
+
         return NodeOutputPointer(
             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:
+            raise InvalidInputReferenceError(
+                message=f"Inputs class '{value.inputs_class.__qualname__}' has no attribute '{value.name}'",
+                inputs_class_name=value.inputs_class.__qualname__,
+                attribute_name=value.name,
+            )
         workflow_input_display = display_context.global_workflow_input_displays[value]
         return InputVariablePointer(data=InputVariableData(input_variable_id=str(workflow_input_display.id)))
     if isinstance(value, VellumSecretReference):
@@ -137,6 +206,26 @@ def create_node_input_value_pointer_rule(
         return ExecutionCounterPointer(
             data=ExecutionCounterData(node_id=str(node_class_display.node_id)),
         )
+    if isinstance(value, EnvironmentVariableReference):
+        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)

vellum_ee/workflows/display/vellum.py

@@ -1,135 +1,9 @@
-from dataclasses import dataclass
-from uuid import UUID
-from typing import List, Literal, Optional
+from typing import List, Literal
 
-from vellum.core import UniversalBaseModel
-from vellum_ee.workflows.display.base import (
-    EdgeDisplayOverrides,
-    EntrypointDisplayOverrides,
-    StateValueDisplayOverrides,
-    WorkflowInputsDisplay,
-    WorkflowMetaDisplayOverrides,
-    WorkflowOutputDisplayOverrides,
-)
-from vellum_ee.workflows.display.base import WorkflowDisplayData  # noqa: F401 - Remove in 0.15.0
-from vellum_ee.workflows.display.base import WorkflowDisplayDataViewport  # noqa: F401 - Remove in 0.15.0
-from vellum_ee.workflows.display.editor.types import NodeDisplayComment  # noqa: F401 - Remove in 0.15.0
-from vellum_ee.workflows.display.editor.types import NodeDisplayData
-from vellum_ee.workflows.display.editor.types import NodeDisplayPosition  # noqa: F401 - Remove in 0.15.0
+from vellum.client.core import UniversalBaseModel
 from vellum_ee.workflows.display.utils.vellum import NodeInputValuePointerRule
 
 
-@dataclass
-class WorkflowMetaVellumDisplayOverrides(WorkflowMetaDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowMetaDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class WorkflowMetaVellumDisplay(WorkflowMetaVellumDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowMetaDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class WorkflowInputsVellumDisplayOverrides(WorkflowInputsDisplay):
-    """
-    DEPRECATED: Use WorkflowInputsDisplay instead. Will be removed in 0.15.0
-    """
-
-    required: Optional[bool] = None
-
-
-@dataclass
-class WorkflowInputsVellumDisplay(WorkflowInputsVellumDisplayOverrides):
-    pass
-
-
-@dataclass
-class StateValueVellumDisplayOverrides(StateValueDisplayOverrides):
-    """
-    DEPRECATED: Use StateValueDisplay instead. Will be removed in 0.15.0
-    """
-
-    required: Optional[bool] = None
-
-
-@dataclass
-class StateValueVellumDisplay(StateValueVellumDisplayOverrides):
-    """
-    DEPRECATED: Use StateValueDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class EdgeVellumDisplayOverrides(EdgeDisplayOverrides):
-    """
-    DEPRECATED: Use EdgeDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class EdgeVellumDisplay(EdgeVellumDisplayOverrides):
-    """
-    DEPRECATED: Use EdgeDisplay instead. Will be removed in 0.15.0
-    """
-
-    source_node_id: UUID
-    source_handle_id: UUID
-    target_node_id: UUID
-    target_handle_id: UUID
-    type: Literal["DEFAULT"] = "DEFAULT"
-
-
-@dataclass
-class EntrypointVellumDisplayOverrides(EntrypointDisplayOverrides):
-    """
-    DEPRECATED: Use EntrypointDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class EntrypointVellumDisplay(EntrypointVellumDisplayOverrides):
-    """
-    DEPRECATED: Use EntrypointDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class WorkflowOutputVellumDisplayOverrides(WorkflowOutputDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowOutputDisplay instead. Will be removed in 0.15.0
-    """
-
-    label: Optional[str] = None
-    node_id: Optional[UUID] = None
-    display_data: Optional[NodeDisplayData] = None
-    target_handle_id: Optional[UUID] = None
-
-
-@dataclass
-class WorkflowOutputVellumDisplay(WorkflowOutputVellumDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowOutputDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
 class NodeInputValuePointer(UniversalBaseModel):
     rules: List[NodeInputValuePointerRule]
     combinator: Literal["OR"] = "OR"

vellum_ee/workflows/display/workflows/__init__.py

@@ -2,4 +2,3 @@
 
 # These must be imported to register the classes in the registry
 from .base_workflow_display import BaseWorkflowDisplay
-from .vellum_workflow_display import VellumWorkflowDisplay