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

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/client/types/integration_name.py

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

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/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/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/workflows/nodes/displayable/set_state_node/__init__.py

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

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)

vellum/workflows/nodes/displayable/set_state_node/tests/test_node.py

@@ -0,0 +1,212 @@
+import pytest
+from typing import List
+
+from vellum import ChatMessage
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.displayable.set_state_node import SetStateNode
+from vellum.workflows.state.base import BaseState
+
+
+def test_set_state_node_with_single_value():
+    """Test that SetStateNode can set a single state value."""
+
+    # GIVEN a state with a string field
+    class TestState(BaseState):
+        message: str = ""
+
+    # AND a SetStateNode that sets a message
+    class SetMessage(SetStateNode[TestState]):
+        operations = {"message": "Hello, World!"}
+
+    # WHEN we run the node
+    state = TestState()
+    node = SetMessage(state=state)
+    outputs = node.run()
+
+    # THEN the result should be the set value
+    assert outputs.result == {"message": "Hello, World!"}
+
+    # AND the state should be updated
+    assert state.message == "Hello, World!"
+
+
+def test_set_state_node_with_multiple_values():
+    """Test that SetStateNode can set multiple state values at once with expressions."""
+
+    # GIVEN a state with multiple fields
+    class TestState(BaseState):
+        counter: int = 0
+        total_tokens: int = 0
+        user_tokens: int = 5
+        assistant_tokens: int = 10
+
+    # AND a SetStateNode that sets multiple values with expressions using add
+    class UpdateMultipleState(SetStateNode[TestState]):
+        operations = {
+            "counter": 1,
+            "total_tokens": TestState.user_tokens + TestState.assistant_tokens,
+        }
+
+    # WHEN we run the node
+    state = TestState()
+    node = UpdateMultipleState(state=state)
+    outputs = node.run()
+
+    # THEN all state values should be updated
+    assert state.counter == 1
+    assert state.total_tokens == 15
+
+    # AND the result should contain all updates
+    assert outputs.result == {"counter": 1, "total_tokens": 15}
+
+
+def test_set_state_node_with_chat_history_and_concat():
+    """Test that SetStateNode can concatenate chat history using concat."""
+
+    # GIVEN a state with chat history
+    class TestState(BaseState):
+        chat_history: List[ChatMessage] = []
+
+    # AND initial chat history in state
+    initial_messages = [ChatMessage(role="ASSISTANT", text="Hi there")]
+
+    # AND a SetStateNode that uses concat() method to add to chat history
+    class UpdateChatHistory(SetStateNode[TestState]):
+        operations = {"chat_history": TestState.chat_history.concat([ChatMessage(role="USER", text="Hello")])}
+
+    # WHEN we run the node with initial state
+    state = TestState(chat_history=initial_messages)
+
+    node = UpdateChatHistory(state=state)
+    outputs = node.run()
+
+    # THEN the state should have concatenated chat history
+    assert len(state.chat_history) == 2
+    assert state.chat_history[0].role == "ASSISTANT"
+    assert state.chat_history[1].role == "USER"
+    assert state.chat_history[1].text == "Hello"
+
+    # AND the result should contain the updated chat history
+    assert outputs.result == {
+        "chat_history": [ChatMessage(role="ASSISTANT", text="Hi there"), ChatMessage(role="USER", text="Hello")]
+    }
+
+
+def test_set_state_node_with_empty_dict():
+    """Test that SetStateNode handles empty operations gracefully."""
+
+    # GIVEN a state
+    class TestState(BaseState):
+        counter: int = 0
+
+    # AND a SetStateNode with empty operations
+    class EmptySetState(SetStateNode[TestState]):
+        operations = {}
+
+    # WHEN we run the node
+    state = TestState()
+    node = EmptySetState(state=state)
+    outputs = node.run()
+
+    # THEN the result should be an empty dict
+    assert outputs.result == {}
+
+    # AND state should be unchanged
+    assert state.counter == 0
+
+
+def test_set_state_node_modifies_existing_values():
+    """Test that SetStateNode can modify existing state values."""
+
+    # GIVEN a state with existing values
+    class TestState(BaseState):
+        counter: int = 10
+        message: str = "old"
+
+    # AND a SetStateNode that updates existing values
+    class UpdateState(SetStateNode[TestState]):
+        operations = {"counter": 20, "message": "new"}
+
+    # WHEN we run the node
+    state = TestState()
+    node = UpdateState(state=state)
+    outputs = node.run()
+
+    # THEN state should be updated
+    assert state.counter == 20
+    assert state.message == "new"
+
+    # AND result should contain the new values with the correct structure
+    assert outputs.result == {"counter": 20, "message": "new"}
+
+
+def test_set_state_not_existing_value():
+    """Test that SetStateNode raises an error when trying to set a value that doesn't exist in the state."""
+
+    # GIVEN a state
+    class TestState(BaseState):
+        counter: int = 0
+
+    # AND a SetStateNode that tries to set a value that doesn't exist in the state
+    class SetState(SetStateNode[TestState]):
+        operations = {"unknown_value": 10}
+
+    # WHEN we run the node
+    state = TestState()
+    node = SetState(state=state)
+
+    # THEN it should raise a NodeException
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # AND the error should be INVALID_STATE and the message should mention the non-existent attribute
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_STATE
+    assert (
+        "State does not have attribute 'unknown_value'. Only existing state attributes can be set via SetStateNode."
+        == str(exc_info.value)
+    )
+
+
+def test_set_state_node_atomic_order_independent_resolution():
+    """Both operations resolve against the original state before any mutation."""
+
+    class TestState(BaseState):
+        a: int = 1
+        b: int = 0
+
+    class UpdateAB(SetStateNode[TestState]):
+        # Both should see original a=1, so both resolve to 2
+        operations = {
+            "a": TestState.a + 1,
+            "b": TestState.a + 1,
+        }
+
+    state = TestState()
+    node = UpdateAB(state=state)
+    outputs = node.run()
+
+    assert state.a == 2
+    assert state.b == 2
+    assert outputs.result == {"a": 2, "b": 2}
+
+
+def test_set_state_node_no_partial_update_on_error():
+    """If a later operation is invalid, no earlier changes should be applied."""
+
+    class TestState(BaseState):
+        a: int = 5
+
+    class PartialFail(SetStateNode[TestState]):
+        operations = {
+            "a": 42,  # would change a
+            "missing": 1,  # invalid attribute triggers NodeException
+        }
+
+    state = TestState()
+    node = PartialFail(state=state)
+
+    with pytest.raises(NodeException):
+        node.run()
+
+    assert state.a == 5

vellum/workflows/sandbox.py

@@ -1,4 +1,4 @@
-from typing import Generic, Optional, Sequence, Union
+from typing import Any, Dict, Generic, Optional, Sequence, Union
 
 import dotenv
 
@@ -51,11 +51,21 @@ class WorkflowSandboxRunner(Generic[WorkflowType]):
 
         selected_inputs = self._inputs[index]
 
+        raw_inputs: Union[BaseInputs, Dict[str, Any]]
         if isinstance(selected_inputs, DatasetRow):
-            selected_inputs = selected_inputs.inputs
+            raw_inputs = selected_inputs.inputs
+        else:
+            raw_inputs = selected_inputs
+
+        inputs_for_stream: BaseInputs
+        if isinstance(raw_inputs, dict):
+            inputs_class = type(self._workflow).get_inputs_class()
+            inputs_for_stream = inputs_class(**raw_inputs)
+        else:
+            inputs_for_stream = raw_inputs
 
         events = self._workflow.stream(
-            inputs=selected_inputs,
+            inputs=inputs_for_stream,
             event_filter=root_workflow_event_filter,
         )
 

vellum/workflows/tests/test_dataset_row.py

@@ -126,3 +126,23 @@ def test_dataset_row_with_empty_inputs():
 
     assert serialized_dict["label"] == "test_label"
     assert serialized_dict["inputs"] == {}
+
+
+def test_dataset_row_with_dict_inputs():
+    """
+    Test that DatasetRow can accept a dict for the inputs parameter.
+    """
+
+    # GIVEN a dict with input data
+    inputs_dict = {"message": "Hello World", "count": 42}
+
+    dataset_row = DatasetRow(label="test_label", inputs=inputs_dict)
+
+    # THEN the inputs should remain as a dict
+    assert isinstance(dataset_row.inputs, dict)
+
+    # AND the serialized dict should contain the correct data
+    serialized_dict = dataset_row.model_dump()
+    assert serialized_dict["label"] == "test_label"
+    assert serialized_dict["inputs"]["message"] == "Hello World"
+    assert serialized_dict["inputs"]["count"] == 42

vellum/workflows/tests/test_sandbox.py

@@ -2,6 +2,7 @@ import pytest
 from typing import List
 
 from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.inputs.dataset_row import DatasetRow
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.sandbox import WorkflowSandboxRunner
 from vellum.workflows.state.base import BaseState
@@ -60,3 +61,42 @@ def test_sandbox_runner__happy_path(mock_logger, run_kwargs, expected_last_log):
         "----------------------------------",
         expected_last_log,
     ]
+
+
+def test_sandbox_runner_with_dict_inputs(mock_logger):
+    """
+    Test that WorkflowSandboxRunner can run with dict inputs in DatasetRow.
+    """
+
+    # GIVEN we capture the logs to stdout
+    logs = []
+    mock_logger.return_value.info.side_effect = lambda msg: logs.append(msg)
+
+    class Inputs(BaseInputs):
+        message: str
+
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result = Inputs.message
+
+    class Workflow(BaseWorkflow[Inputs, BaseState]):
+        graph = StartNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_output = StartNode.Outputs.result
+
+    dataset = [
+        DatasetRow(label="test_row", inputs={"message": "Hello from dict"}),
+    ]
+
+    # WHEN we run the sandbox with the DatasetRow containing dict inputs
+    runner = WorkflowSandboxRunner(workflow=Workflow(), dataset=dataset)
+    runner.run()
+
+    assert logs == [
+        "Just started Node: StartNode",
+        "Just finished Node: StartNode",
+        "Workflow fulfilled!",
+        "----------------------------------",
+        "final_output: Hello from dict",
+    ]

vellum/workflows/tests/triggers/{test_vellum_integration_trigger.py → test_integration_trigger.py}

@@ -1,10 +1,10 @@
-"""Tests for VellumIntegrationTrigger."""
+"""Tests for IntegrationTrigger."""
 
 import pytest
 
 from vellum.workflows.constants import VellumIntegrationProviderType
 from vellum.workflows.references.trigger import TriggerAttributeReference
-from vellum.workflows.triggers.vellum_integration import VellumIntegrationTrigger
+from vellum.workflows.triggers.integration import IntegrationTrigger
 
 
 def test_requires_config():
@@ -13,27 +13,27 @@ def test_requires_config():
     # This should fail - no Config class
     with pytest.raises(TypeError, match="Config"):
 
-        class BadTrigger1(VellumIntegrationTrigger):
+        class BadTrigger1(IntegrationTrigger):
             message: str
             # Missing Config!
 
     # This should fail - incomplete Config
     with pytest.raises(TypeError, match="provider"):
 
-        class BadTrigger2(VellumIntegrationTrigger):
+        class BadTrigger2(IntegrationTrigger):
             message: str
 
-            class Config(VellumIntegrationTrigger.Config):
+            class Config(IntegrationTrigger.Config):
                 integration_name = "SLACK"
                 slug = "slack_new_message"
                 # Missing provider!
 
     # This should work
-    class GoodTrigger(VellumIntegrationTrigger):
+    class GoodTrigger(IntegrationTrigger):
         message: str
         user: str
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "SLACK"
             slug = "slack_new_message"
@@ -44,12 +44,12 @@ def test_requires_config():
 def test_top_level_annotations_create_references():
     """Top-level type annotations (webhook event attributes) automatically create TriggerAttributeReference."""
 
-    class SlackTrigger(VellumIntegrationTrigger):
+    class SlackTrigger(IntegrationTrigger):
         message: str
         user: str
         timestamp: float
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "SLACK"
             slug = "slack_new_message"
@@ -69,18 +69,18 @@ def test_top_level_annotations_create_references():
 def test_attribute_ids_include_class_name():
     """Attribute IDs should include class name (like nodes)."""
 
-    class Trigger1(VellumIntegrationTrigger):
+    class Trigger1(IntegrationTrigger):
         message: str
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "SLACK"
             slug = "slack_new_message"
 
-    class Trigger2(VellumIntegrationTrigger):
+    class Trigger2(IntegrationTrigger):
         message: str
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "SLACK"
             slug = "slack_new_message"
@@ -93,12 +93,12 @@ def test_attribute_ids_include_class_name():
 def test_populates_dynamic_attributes():
     """Trigger dynamically populates attributes from event_data keys."""
 
-    class GithubPushTrigger(VellumIntegrationTrigger):
+    class GithubPushTrigger(IntegrationTrigger):
         repository: str
         branch: str
         commits: list
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "GITHUB"
             slug = "github_push_event"
@@ -109,7 +109,7 @@ def test_populates_dynamic_attributes():
         "commits": ["abc123", "def456"],
     }
 
-    trigger = GithubPushTrigger(event_data=event_data)
+    trigger = GithubPushTrigger(**event_data)
 
     assert trigger.repository == "vellum-ai/workflows"
     assert trigger.branch == "main"
@@ -119,17 +119,17 @@ def test_populates_dynamic_attributes():
 def test_to_trigger_attribute_values():
     """to_trigger_attribute_values returns correct attribute mappings."""
 
-    class SlackTrigger(VellumIntegrationTrigger):
+    class SlackTrigger(IntegrationTrigger):
         message: str
         channel: str
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "SLACK"
             slug = "slack_new_message"
 
     event_data = {"message": "Hello", "channel": "C123"}
-    trigger = SlackTrigger(event_data=event_data)
+    trigger = SlackTrigger(**event_data)
 
     attr_values = trigger.to_trigger_attribute_values()
 
@@ -142,15 +142,15 @@ def test_to_trigger_attribute_values():
 def test_empty_event_data():
     """Trigger handles empty event data gracefully."""
 
-    class SlackTrigger(VellumIntegrationTrigger):
+    class SlackTrigger(IntegrationTrigger):
         message: str
 
-        class Config(VellumIntegrationTrigger.Config):
+        class Config(IntegrationTrigger.Config):
             provider = VellumIntegrationProviderType.COMPOSIO
             integration_name = "SLACK"
             slug = "slack_new_message"
 
-    trigger = SlackTrigger(event_data={})
+    trigger = SlackTrigger()
 
     attr_values = trigger.to_trigger_attribute_values()
     assert attr_values == {}

vellum/workflows/triggers/__init__.py

@@ -2,6 +2,5 @@ from vellum.workflows.triggers.base import BaseTrigger
 from vellum.workflows.triggers.integration import IntegrationTrigger
 from vellum.workflows.triggers.manual import ManualTrigger
 from vellum.workflows.triggers.schedule import ScheduleTrigger
-from vellum.workflows.triggers.vellum_integration import VellumIntegrationTrigger
 
-__all__ = ["BaseTrigger", "IntegrationTrigger", "ManualTrigger", "ScheduleTrigger", "VellumIntegrationTrigger"]
+__all__ = ["BaseTrigger", "IntegrationTrigger", "ManualTrigger", "ScheduleTrigger"]