vellum-ai 1.8.6__tar.gz → 1.9.0__tar.gz

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.8.6
+Version: 1.9.0
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

vellum_ai-1.8.6/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_vellum_integration_trigger_serialization.py → vellum_ai-1.9.0/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_integration_trigger_serialization.py

@@ -1,18 +1,18 @@
-"""Tests for VellumIntegrationTrigger serialization."""
+"""Tests for IntegrationTrigger serialization."""
 
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.triggers.vellum_integration import VellumIntegrationTrigger
+from vellum.workflows.triggers.integration import IntegrationTrigger
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
 def test_vellum_integration_trigger_serialization():
-    """VellumIntegrationTrigger subclass serializes with class name and module path."""
+    """IntegrationTrigger subclass serializes with class name and module path."""
 
-    # Create a custom VellumIntegrationTrigger subclass
-    class SlackMessageTrigger(VellumIntegrationTrigger):
+    # Create a custom IntegrationTrigger subclass
+    class SlackMessageTrigger(IntegrationTrigger):
         """Custom Slack message trigger."""
 
         message: str
@@ -71,7 +71,7 @@ def test_vellum_integration_trigger_serialization():
 def test_vellum_integration_trigger_id_consistency():
     """Validates trigger and attribute IDs match between definitions and references."""
 
-    class SlackMessageTrigger(VellumIntegrationTrigger):
+    class SlackMessageTrigger(IntegrationTrigger):
         message: str
         channel: str
 
@@ -147,7 +147,7 @@ def test_vellum_integration_trigger_id_consistency():
 def test_trigger_module_paths_are_canonical():
     """Validates trigger module_path and class_name for consistent codegen."""
 
-    class TestSlackTrigger(VellumIntegrationTrigger):
+    class TestSlackTrigger(IntegrationTrigger):
         message: str
 
         class Config:
@@ -184,7 +184,7 @@ def test_trigger_module_paths_are_canonical():
 def test_integration_trigger_no_entrypoint_node():
     """IntegrationTrigger-only workflows use trigger ID in edges."""
 
-    class SlackMessageTrigger(VellumIntegrationTrigger):
+    class SlackMessageTrigger(IntegrationTrigger):
         message: str
 
         class Config:

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -25,8 +25,8 @@ from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port,
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.state.encoder import DefaultStateEncoder
+from vellum.workflows.triggers.integration import IntegrationTrigger
 from vellum.workflows.triggers.manual import ManualTrigger
-from vellum.workflows.triggers.vellum_integration import VellumIntegrationTrigger
 from vellum.workflows.types.core import Json, JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.types.utils import get_original_base
@@ -190,7 +190,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         # For IntegrationTrigger only: no ENTRYPOINT node (use trigger ID directly in edges)
         manual_trigger_edges = [edge for edge in trigger_edges if issubclass(edge.trigger_class, ManualTrigger)]
         integration_trigger_edges = [
-            edge for edge in trigger_edges if issubclass(edge.trigger_class, VellumIntegrationTrigger)
+            edge for edge in trigger_edges if issubclass(edge.trigger_class, IntegrationTrigger)
         ]
 
         has_manual_trigger = len(manual_trigger_edges) > 0
@@ -536,12 +536,12 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                     f"{edge.trigger_class.__name__} in the same workflow."
                 )
 
-        # Get the trigger type from the mapping, or check if it's a VellumIntegrationTrigger subclass
+        # Get the trigger type from the mapping, or check if it's an IntegrationTrigger subclass
         trigger_type = trigger_type_mapping.get(trigger_class)
         if trigger_type is None:
-            # Check if it's a VellumIntegrationTrigger subclass
+            # Check if it's an IntegrationTrigger subclass
 
-            if issubclass(trigger_class, VellumIntegrationTrigger):
+            if issubclass(trigger_class, IntegrationTrigger):
                 trigger_type = WorkflowTriggerType.INTEGRATION
             else:
                 raise ValueError(

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/ee/vellum_ee/workflows/server/virtual_file_loader.py

@@ -85,13 +85,11 @@ class VirtualFileLoader(importlib.abc.Loader):
 
     def _is_package_directory(self, fullname: str) -> bool:
         """Check if directory contains .py files that should be treated as a package."""
-        directory_prefix = fullname.replace(".", "/") + "/"
-
-        # Exclude top-level display directory from auto-generation as it typically has
-        # specific __init__.py content that shouldn't be replaced with empty files.
-        if directory_prefix == "display/":
-            return False
+        if fullname == "":
+            # This is the root namespace, so it's a package directory
+            return True
 
+        directory_prefix = fullname.replace(".", "/") + "/"
         for file_path in self.files.keys():
             if file_path.startswith(directory_prefix):
                 if file_path.endswith(".py") and not file_path.endswith("__init__.py"):
@@ -104,10 +102,40 @@ class VirtualFileLoader(importlib.abc.Loader):
 
     def _generate_init_content(self, fullname: str) -> tuple[str, str]:
         """Auto-generate empty __init__.py content to mark directory as a package."""
+        if fullname == "":
+            if any(file_path.startswith("display/") for file_path in self.files.keys()):
+                return (
+                    "__init__.py",
+                    """\
+from .display import *
+""",
+                )
+            else:
+                return "__init__.py", ""
+
         directory_prefix = fullname.replace(".", "/") + "/"
         file_path = directory_prefix + "__init__.py"
 
-        code = ""
+        if not fullname.startswith("display"):
+            return "__init__.py", ""
+
+        # We have to go through the effort of exhaustively importing all display modules, because that's how
+        # dynamically annotate our Workflow graph.
+        all_nested_files = [
+            file_path[len(directory_prefix) :]
+            for file_path in self.files.keys()
+            if file_path.startswith(directory_prefix)
+        ]
+
+        top_level_items = {file_path.split("/")[0] for file_path in all_nested_files}
+
+        top_level_modules = [
+            re.sub(r"\.py$", "", file_path)
+            for file_path in top_level_items
+            if file_path.endswith(".py") or "." not in file_path
+        ]
+
+        code = "\n".join([f"from .{file_path} import *" for file_path in top_level_modules])
 
         return file_path, code
 
@@ -126,7 +154,7 @@ class VirtualFileFinder(BaseWorkflowFinder):
 
     def find_spec(self, fullname: str, path, target=None):
         module_info = self.loader._resolve_module(fullname)
-        if module_info:
+        if isinstance(module_info, tuple) and len(module_info) == 2:
             file_path, _ = module_info
             is_package = file_path.endswith("__init__.py")
             return importlib.machinery.ModuleSpec(

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/ee/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,7 +29,6 @@ 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
@@ -902,3 +903,81 @@ class TestSubworkflowDeploymentNode(SubworkflowDeploymentNode):
 
     # 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_ai-1.8.6 → vellum_ai-1.9.0}/ee/vellum_ee/workflows/tests/test_virtual_files.py

@@ -121,12 +121,9 @@ class WorkflowDisplay(BaseWorkflowDisplay):
 
         spec = importlib.util.find_spec(f"{namespace}.display")
 
-        # THEN the spec should be None because we don't want to auto-generate display/__init__.py
-        # If the spec exists, it means an empty __init__.py was auto-generated (BAD)
-        assert spec is None, (
-            "display directory should NOT have auto-generated __init__.py. "
-            "Display directories require specific __init__.py content that shouldn't be empty."
-        )
+        # 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

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "1.8.6"
+version = "1.9.0"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/src/vellum/client/core/client_wrapper.py

@@ -27,10 +27,10 @@ class BaseClientWrapper:
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
-            "User-Agent": "vellum-ai/1.8.6",
+            "User-Agent": "vellum-ai/1.9.0",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "1.8.6",
+            "X-Fern-SDK-Version": "1.9.0",
             **(self.get_custom_headers() or {}),
         }
         if self._api_version is not None:

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/src/vellum/client/types/integration_name.py

@@ -26,6 +26,7 @@ IntegrationName = typing.Union[
         "APOLLO",
         "ASANA",
         "ATLASSIAN",
+        "BITBUCKET",
         "BOX",
         "BREVO",
         "BROWSERBASE_TOOL",

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/src/vellum/workflows/inputs/dataset_row.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 from pydantic import Field, field_serializer
 
@@ -12,25 +12,27 @@ class DatasetRow(UniversalBaseModel):
 
     Attributes:
         label: String label for the dataset row
-        inputs: BaseInputs instance containing the input data
+        inputs: BaseInputs instance or dict containing the input data
     """
 
     label: str
-    inputs: BaseInputs = Field(default_factory=BaseInputs)
+    inputs: Union[BaseInputs, Dict[str, Any]] = Field(default_factory=BaseInputs)
 
     @field_serializer("inputs")
-    def serialize_inputs(self, inputs: BaseInputs) -> Dict[str, Any]:
+    def serialize_inputs(self, inputs: Union[BaseInputs, Dict[str, Any]]) -> Dict[str, Any]:
         """
-        Custom serializer for BaseInputs that converts it to a dictionary.
+        Custom serializer for inputs that converts it to a dictionary.
 
         Args:
-            inputs: BaseInputs instance to serialize
+            inputs: Either a BaseInputs instance or dict to serialize
 
         Returns:
             Dictionary representation of the inputs
         """
-        result = {}
+        if isinstance(inputs, dict):
+            return inputs
 
+        result = {}
         for input_descriptor, value in inputs:
             if not input_descriptor.name.startswith("__"):
                 result[input_descriptor.name] = value

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/src/vellum/workflows/nodes/displayable/final_output_node/node.py

@@ -59,6 +59,10 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
             return
 
         declared_output_type = cls.get_output_type()
+
+        if declared_output_type is Any:
+            return
+
         value_descriptor = None
 
         if "value" in outputs_class.__dict__:

{vellum_ai-1.8.6 → vellum_ai-1.9.0}/src/vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py

@@ -2,10 +2,12 @@ import pytest
 from typing import Any, Dict
 
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.nodes.displayable.inline_prompt_node import InlinePromptNode
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import Json
 
 
 def test_final_output_node__mismatched_output_type_should_raise_exception_when_ran():
@@ -85,3 +87,29 @@ def test_final_output_node__dict_and_Dict_should_be_compatible():
         DictOutputNode.__validate__()
     except ValueError as e:
         pytest.fail(f"Validation should not raise an exception for dict/Dict compatibility: {e}")
+
+
+def test_final_output_node__any_output_type_should_accept_json():
+    """
+    Tests that FinalOutputNode with Any output type accepts a TemplatingNode with Json output type.
+    """
+
+    # GIVEN a TemplatingNode with Json output type
+    class JsonTemplatingNode(TemplatingNode[BaseState, Json]):
+        """Templating node that outputs Json."""
+
+        template = '{"key": "value"}'
+
+    # AND a FinalOutputNode with Any output type referencing the TemplatingNode
+    class AnyOutputNode(FinalOutputNode[BaseState, Any]):
+        """Output with Any type."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = JsonTemplatingNode.Outputs.result
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        AnyOutputNode.__validate__()
+    except ValueError as e:
+        pytest.fail(f"Validation should not raise an exception when Any accepts Json: {e}")

vellum_ai-1.9.0/src/vellum/workflows/nodes/displayable/set_state_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import SetStateNode
+
+__all__ = [
+    "SetStateNode",
+]

vellum_ai-1.9.0/src/vellum/workflows/nodes/displayable/set_state_node/node.py

@@ -0,0 +1,71 @@
+from typing import Any, ClassVar, Dict, Generic
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.utils import resolve_value
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.types.generics import StateType
+
+
+class SetStateNode(BaseNode[StateType], Generic[StateType]):
+    """
+    A node that sets multiple state values at once.
+
+    This node accepts a dictionary mapping state paths to values and sets them all.
+    Values can be actual values or descriptors.
+
+    Example:
+        operations = {
+            "chat_history": State.chat_history.concat(LazyReference("MyAgentNode.Outputs.chat_history")),
+            "counter": State.counter + 1,
+            "total_tokens": State.user_tokens.add(State.assistant_tokens)
+        }
+
+    Args:
+        operations: Dictionary mapping state attribute names to values (descriptors or actual values)
+    """
+
+    # Dictionary mapping state paths to values
+    operations: ClassVar[Dict[str, Any]] = {}
+
+    class Outputs(BaseOutputs):
+        """
+        The outputs of the SetStateNode.
+
+        result: Dict[str, Any] - Dictionary of all state updates
+        """
+
+        result: Dict[str, Any]
+
+    def run(self) -> Outputs:
+        """
+        Run the node and set all the state values.
+        Resolves descriptors to their actual values before setting state.
+        """
+        # First pass: validate and resolve all operations without mutating state
+        resolved_updates: Dict[str, Any] = {}
+        for path, value in self.operations.items():
+            # Validate the state attribute exists prior to any mutation
+            if not hasattr(self.state, path):
+                raise NodeException(
+                    f"State does not have attribute '{path}'. "
+                    f"Only existing state attributes can be set via SetStateNode.",
+                    code=WorkflowErrorCode.INVALID_STATE,
+                )
+
+            # Resolve the value if it's a descriptor against the current (unmodified) state
+            if isinstance(value, BaseDescriptor):
+                resolved_value = resolve_value(value, self.state)
+            else:
+                resolved_value = value
+
+            resolved_updates[path] = resolved_value
+
+        # Second pass: apply the resolved updates to the state atomically
+        with self.state.__atomic__():
+            for path, resolved_value in resolved_updates.items():
+                setattr(self.state, path, resolved_value)
+
+        return self.Outputs(result=resolved_updates)