vellum-ai 1.8.5__py3-none-any.whl → 1.9.0__py3-none-any.whl

vellum_ee/workflows/tests/test_server.py

@@ -1,17 +1,19 @@
 import pytest
 import sys
-from uuid import uuid4
+from uuid import UUID, uuid4
 from typing import Type, cast
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.code_executor_response import CodeExecutorResponse
 from vellum.client.types.number_vellum_value import NumberVellumValue
 from vellum.workflows import BaseWorkflow
+from vellum.workflows.events.workflow import WorkflowExecutionInitiatedBody
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.utils.uuids import generate_workflow_deployment_prefix
 from vellum.workflows.utils.zip import zip_file_map
+from vellum_ee.workflows.display.utils.events import event_enricher
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
@@ -27,20 +29,12 @@ def test_load_workflow_event_display_context():
 def test_load_from_module__lazy_reference_in_file_loader():
     # GIVEN a workflow module with a node containing a lazy reference
     files = {
-        "__init__.py": "",
         "workflow.py": """\
 from vellum.workflows import BaseWorkflow
 from .nodes.start_node import StartNode
 
 class Workflow(BaseWorkflow):
     graph = StartNode
-""",
-        "nodes/__init__.py": """\
-from .start_node import StartNode
-
-__all__ = [
-    "StartNode",
-]
 """,
         "nodes/start_node.py": """\
 from vellum.workflows.nodes import BaseNode
@@ -98,11 +92,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = CodeExecutionNode.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/code_execution_node/__init__.py": """
 from typing import Any
@@ -174,11 +163,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = CodeExecutionNode.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -250,11 +234,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = CodeExecutionNode.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -326,11 +305,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = Subworkflow.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .subworkflow import Subworkflow
-
-__all__ = ["Subworkflow"]
 """,
         "nodes/subworkflow/__init__.py": """
 from vellum.workflows.nodes.displayable import InlineSubworkflowNode
@@ -339,11 +313,6 @@ from .workflow import SubworkflowWorkflow
 
 class Subworkflow(InlineSubworkflowNode):
     subworkflow = SubworkflowWorkflow
-""",
-        "nodes/subworkflow/nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/subworkflow/nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -424,11 +393,6 @@ class Workflow(BaseWorkflow):
     
     class Outputs(BaseWorkflow.Outputs):  # noqa: W293
        results = MapNode.Outputs.final_output
-""",
-        "nodes/__init__.py": """
-from .map_node import MapNode
-
-__all__ = ["MapNode"]
 """,
         "nodes/map_node/__init__.py": """
 from vellum.workflows.nodes.core.map_node import MapNode as BaseMapNode
@@ -439,11 +403,6 @@ class MapNode(BaseMapNode):
     items = ["foo", "bar", "baz"]
     subworkflow = MapNodeWorkflow
     max_concurrency = 4
-""",
-        "nodes/map_node/nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/map_node/nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -520,7 +479,6 @@ from .nodes.broken_node import BrokenNode
 class Workflow(BaseWorkflow):
     graph = BrokenNode
 """,
-        "nodes/__init__.py": "",
         "nodes/broken_node.py": """\
 from vellum.workflows.nodes import BaseNode
 
@@ -559,7 +517,6 @@ from .nodes.broken_node import BrokenNode
 class Workflow(BaseWorkflow):
     graph = BrokenNode
 """,
-        "nodes/__init__.py": "",
         "nodes/broken_node.py": """\
 from vellum.workflows.nodes import BaseNode
 
@@ -819,19 +776,9 @@ class TestSubworkflowDeploymentNode(SubworkflowDeploymentNode):
     files = {
         "__init__.py": "",
         "workflow.py": parent_workflow_code,
-        "nodes/__init__.py": """
-from .subworkflow_deployment_node import TestSubworkflowDeploymentNode
-
-__all__ = ["TestSubworkflowDeploymentNode"]
-""",
         "nodes/subworkflow_deployment_node.py": parent_node_code,
         f"{expected_prefix}/__init__.py": "",
         f"{expected_prefix}/workflow.py": mock_workflow_code,
-        f"{expected_prefix}/nodes/__init__.py": """
-from .test_node import TestNode
-
-__all__ = ["TestNode"]
-""",
         f"{expected_prefix}/nodes/test_node.py": test_node_code,
     }
 
@@ -919,11 +866,6 @@ class TestSubworkflowDeploymentNode(SubworkflowDeploymentNode):
     subworkflow_files = {
         "__init__.py": "",
         "workflow.py": mock_workflow_code,
-        "nodes/__init__.py": """
-from .test_node import TestNode
-
-__all__ = ["TestNode"]
-""",
         "nodes/test_node.py": test_node_code,
     }
 
@@ -931,11 +873,6 @@ __all__ = ["TestNode"]
         "__init__.py": "",
         "inputs.py": inputs_code,
         "workflow.py": parent_workflow_code,
-        "nodes/__init__.py": """
-from .subworkflow_deployment_node import TestSubworkflowDeploymentNode
-
-__all__ = ["TestSubworkflowDeploymentNode"]
-""",
         "nodes/subworkflow_deployment_node.py": parent_node_code,
     }
 
@@ -966,3 +903,81 @@ __all__ = ["TestSubworkflowDeploymentNode"]
 
     # AND the X-Vellum-Always-Success header should be included for graceful error handling
     assert kwargs["request_options"]["additional_headers"]["X-Vellum-Always-Success"] == "true"
+
+
+def test_workflow_initiated_event_includes_display_context_with_output_display_name():
+    """
+    Tests that workflow initiated events include display context with annotated node and output information.
+    """
+    # GIVEN a workflow with a single node that has a single output
+    files = {
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.start_node import StartNode
+
+class Workflow(BaseWorkflow):
+    graph = StartNode
+
+    class Outputs(BaseWorkflow.Outputs):
+        final_result = StartNode.Outputs.result
+""",
+        "nodes/start_node.py": """\
+from vellum.workflows.nodes import BaseNode
+
+class StartNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        result: str = "test output"
+""",
+        "display/nodes/__init__.py": """\
+from .start_node import StartNodeDisplay
+""",
+        "display/nodes/start_node.py": """\
+from uuid import UUID
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
+from ...nodes.start_node import StartNode
+
+class StartNodeDisplay(BaseNodeDisplay[StartNode]):
+    node_id = UUID("11111111-1111-1111-1111-111111111111")
+    output_display = {
+        StartNode.Outputs.result: NodeOutputDisplay(
+            id=UUID("22222222-2222-2222-2222-222222222222"),
+            name="Pretty Result"
+        )
+    }
+""",
+        "display/workflow.py": """\
+""",
+    }
+
+    namespace = str(uuid4())
+    finder = VirtualFileFinder(files, namespace)
+    sys.meta_path.append(finder)
+
+    try:
+        # WHEN we stream the workflow and enrich the initiated event with display context
+        Workflow = BaseWorkflow.load_from_module(namespace)
+        workflow = Workflow()
+        initiated_event = list(workflow.stream())[0]
+        enriched_event = event_enricher(initiated_event)
+
+        # THEN the initiated event should have display context
+        assert isinstance(enriched_event.body, WorkflowExecutionInitiatedBody)
+        body = enriched_event.body
+        assert body.display_context is not None
+
+        # AND the display context should contain the annotated node
+        node_displays = body.display_context.node_displays
+        annotated_node_id = UUID("11111111-1111-1111-1111-111111111111")
+        assert annotated_node_id in node_displays
+
+        # AND the node's output display should contain the output with the display name "Pretty Result"
+        node_display = node_displays[annotated_node_id]
+        assert "result" in node_display.output_display
+
+        # AND the output should map to the annotated output id
+        annotated_output_id = UUID("22222222-2222-2222-2222-222222222222")
+        assert node_display.output_display["result"] == annotated_output_id
+    finally:
+        if finder in sys.meta_path:
+            sys.meta_path.remove(finder)

vellum_ee/workflows/tests/test_virtual_files.py

@@ -80,3 +80,51 @@ def test_base_class_dynamic_import(files):
         assert instance.id
     except Exception as e:
         pytest.fail(f"Failed to create an instance of BaseClass: {e}")
+
+
+def test_display_directory_not_auto_generated():
+    """
+    Test that the top-level display directory is NOT auto-generated with empty __init__.py.
+    Display directories typically have specific __init__.py content (e.g., "from .workflow import *")
+    that should not be replaced with empty auto-generated files.
+    """
+    # GIVEN a workflow with display/workflow.py but NO display/__init__.py
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.bases import BaseNode
+
+class StartNode(BaseNode):
+    pass
+
+class Workflow(BaseWorkflow):
+    graph = StartNode
+""",
+        "display/workflow.py": """\
+from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+
+class WorkflowDisplay(BaseWorkflowDisplay):
+    pass
+""",
+        # Note: NO "display/__init__.py" in files dict
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    sys.meta_path.append(VirtualFileFinder(files, namespace))
+
+    try:
+        # WHEN we try to resolve display/__init__.py
+        import importlib.util
+
+        spec = importlib.util.find_spec(f"{namespace}.display")
+
+        # THEN the spec should be found because we now support dynamic display module imports
+        assert spec
+        assert spec.origin == "display/__init__.py"
+
+    finally:
+        # Clean up
+        sys.meta_path = [finder for finder in sys.meta_path if not isinstance(finder, VirtualFileFinder)]

vellum/workflows/triggers/vellum_integration.py

@@ -1,189 +0,0 @@
-from typing import Any, ClassVar, Dict
-
-from vellum.workflows.constants import VellumIntegrationProviderType
-from vellum.workflows.references.trigger import TriggerAttributeReference
-from vellum.workflows.triggers.base import BaseTriggerMeta
-from vellum.workflows.triggers.integration import IntegrationTrigger
-
-
-class VellumIntegrationTriggerMeta(BaseTriggerMeta):
-    """
-    Custom metaclass for VellumIntegrationTrigger.
-
-    This metaclass extends BaseTriggerMeta to automatically convert type annotations
-    into TriggerAttributeReference objects during class creation. This enables trigger
-    attributes to be referenced in workflow graphs while maintaining type safety.
-    """
-
-    def __new__(mcs, name: str, bases: tuple, namespace: dict, **kwargs: Any) -> "VellumIntegrationTriggerMeta":
-        """Create a new trigger class and set up attribute references."""
-        cls = super().__new__(mcs, name, bases, namespace, **kwargs)
-
-        # Process __annotations__ to create TriggerAttributeReference for each attribute
-        # Only process if class has Config and annotations
-        has_config = hasattr(cls, "Config") and "Config" in namespace
-        if has_config and hasattr(cls, "__annotations__"):
-            # Create TriggerAttributeReference for each annotated attribute
-            for attr_name, attr_type in cls.__annotations__.items():
-                # Skip special attributes and Config
-                if attr_name.startswith("_") or attr_name == "Config":
-                    continue
-
-                # Create reference with proper type
-                reference = TriggerAttributeReference(
-                    name=attr_name, types=(attr_type,), instance=None, trigger_class=cls
-                )
-                # Set as class attribute so it's directly accessible
-                setattr(cls, attr_name, reference)
-
-        return cls
-
-
-class VellumIntegrationTrigger(IntegrationTrigger, metaclass=VellumIntegrationTriggerMeta):
-    """
-    Base class for Vellum-managed integration triggers.
-
-    Subclasses define two types of attributes:
-    1. **Config class**: Specifies how the trigger is configured (provider, integration_name, slug)
-       - These are configuration details users shouldn't need to interact with directly
-    2. **Top-level type annotations**: Define the webhook event payload structure (message, user, channel, etc.)
-       - These become TriggerAttributeReference that can be referenced in workflow nodes
-
-    Examples:
-        Create a Slack trigger:
-            >>> class SlackNewMessageTrigger(VellumIntegrationTrigger):
-            ...     # Event attributes (webhook payload structure)
-            ...     message: str
-            ...     user: str
-            ...     channel: str
-            ...     timestamp: float
-            ...
-            ...     # Configuration (how trigger is set up)
-            ...     class Config(VellumIntegrationTrigger.Config):
-            ...         provider = VellumIntegrationProviderType.COMPOSIO
-            ...         integration_name = "SLACK"
-            ...         slug = "slack_new_message"
-
-        Use in workflow graph:
-            >>> class MyWorkflow(BaseWorkflow):
-            ...     graph = SlackNewMessageTrigger >> ProcessMessageNode
-
-        Reference trigger attributes in nodes:
-            >>> class ProcessNode(BaseNode):
-            ...     class Outputs(BaseNode.Outputs):
-            ...         text = SlackNewMessageTrigger.message
-            ...         channel = SlackNewMessageTrigger.channel
-
-        Instantiate for testing:
-            >>> trigger = SlackNewMessageTrigger(event_data={
-            ...     "message": "Hello world",
-            ...     "channel": "C123456",
-            ...     "user": "U123",
-            ...     "timestamp": 1234567890.0,
-            ... })
-            >>> trigger.message
-            'Hello world'
-    """
-
-    class Config:
-        """
-        Configuration for VellumIntegrationTrigger subclasses.
-
-        Defines how the trigger connects to the integration provider. These settings
-        specify which integration and which specific trigger type to use.
-        """
-
-        provider: ClassVar[VellumIntegrationProviderType]
-        integration_name: ClassVar[str]
-        slug: ClassVar[str]
-
-    def __init_subclass__(cls, **kwargs: Any) -> None:
-        """Validate that subclasses define required Config class with all required fields."""
-        super().__init_subclass__(**kwargs)
-
-        # Skip validation for the base class itself
-        if cls.__name__ == "VellumIntegrationTrigger":
-            return
-
-        # Require Config class with required fields
-        if not hasattr(cls, "Config") or cls.Config is VellumIntegrationTrigger.Config:
-            raise TypeError(
-                f"{cls.__name__} must define a nested Config class. "
-                f"Example:\n"
-                f"  class {cls.__name__}(VellumIntegrationTrigger):\n"
-                f"      message: str\n"
-                f"      class Config(VellumIntegrationTrigger.Config):\n"
-                f"          provider = VellumIntegrationProviderType.COMPOSIO\n"
-                f"          integration_name = 'SLACK'\n"
-                f"          slug = 'slack_new_message'"
-            )
-
-        # Validate Config class has required fields
-        config_cls = cls.Config
-        required_fields = ["provider", "integration_name", "slug"]
-        for field in required_fields:
-            if not hasattr(config_cls, field):
-                raise TypeError(
-                    f"{cls.__name__}.Config must define '{field}'. " f"Required fields: {', '.join(required_fields)}"
-                )
-
-    def __init__(self, event_data: dict):
-        """
-        Initialize trigger with event data from the integration.
-
-        The trigger dynamically populates its attributes based on the event_data
-        dictionary keys. Any key in event_data becomes an accessible attribute.
-
-        Args:
-            event_data: Raw event data from the integration. Keys become trigger attributes.
-
-        Examples:
-            >>> class SlackTrigger(VellumIntegrationTrigger):
-            ...     message: str
-            ...     channel: str
-            ...     user: str
-            ...
-            ...     class Config(VellumIntegrationTrigger.Config):
-            ...         provider = VellumIntegrationProviderType.COMPOSIO
-            ...         integration_name = "SLACK"
-            ...         slug = "slack_new_message"
-            >>> trigger = SlackTrigger(event_data={
-            ...     "message": "Hello",
-            ...     "channel": "C123",
-            ...     "user": "U456"
-            ... })
-            >>> trigger.message
-            'Hello'
-            >>> trigger.channel
-            'C123'
-        """
-        super().__init__(event_data)
-
-        # Dynamically populate instance attributes from event_data.
-        # This allows any key in event_data to become an accessible attribute:
-        # event_data={"message": "Hi"} → trigger.message == "Hi"
-        for key, value in event_data.items():
-            setattr(self, key, value)
-
-    def to_trigger_attribute_values(self) -> Dict["TriggerAttributeReference[Any]", Any]:
-        """
-        Materialize attribute descriptor/value pairs for this trigger instance.
-
-        For VellumIntegrationTrigger, this includes all dynamic attributes from event_data.
-        """
-        attribute_values: Dict["TriggerAttributeReference[Any]", Any] = {}
-
-        # Unlike the base class which iterates over type(self) (predefined annotations),
-        # we iterate over event_data keys since our attributes are discovered dynamically
-        # from the actual event data received during workflow execution.
-        # The base class approach: for reference in type(self)
-        # Our approach: for attr_name in self._event_data.keys()
-        for attr_name in self._event_data.keys():
-            # Get the class-level reference for this attribute (created by __new__ from annotations)
-            # Unknown keys can appear in webhook payloads, so gracefully skip them if the
-            # trigger class doesn't expose a corresponding reference.
-            reference = getattr(type(self), attr_name, None)
-            if isinstance(reference, TriggerAttributeReference):
-                attribute_values[reference] = getattr(self, attr_name)
-
-        return attribute_values