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

vellum/workflows/state/tests/test_state.py

@@ -2,8 +2,12 @@ import pytest
 from copy import deepcopy
 import json
 from queue import Queue
-from typing import Dict, List, cast
+import threading
+from typing import Any, Dict, List, Optional, cast
 
+from pydantic import Field
+
+from vellum import ChatMessage
 from vellum.utils.json_encoder import VellumJsonEncoder
 from vellum.workflows.constants import undefined
 from vellum.workflows.nodes.bases import BaseNode
@@ -243,3 +247,267 @@ def test_state_deepcopy_handles_undefined_values():
 
     # THEN the undefined values are preserved
     assert deepcopied_state.meta.node_outputs[MockNode.Outputs.baz] == {"foo": undefined}
+
+
+def test_base_state_initializes_field_with_default_factory():
+    """Test that BaseState properly initializes fields with Field(default_factory=...)."""
+
+    # GIVEN a state class with fields using Field(default_factory=...)
+    class TestState(BaseState):
+        chat_history: List[str] = Field(default_factory=list)
+        items: Dict[str, int] = Field(default_factory=dict)
+        counter: int = Field(default_factory=lambda: 0)
+
+    # WHEN we create a state instance without providing values
+    state = TestState()
+
+    # THEN the fields should be initialized with the factory results, not FieldInfo objects
+    assert isinstance(state.chat_history, list)
+    assert state.chat_history == []
+    assert isinstance(state.items, dict)
+    assert state.items == {}
+    assert isinstance(state.counter, int)
+    assert state.counter == 0
+
+    # AND we should be able to modify them
+    state.chat_history.append("message1")
+    state.items["key1"] = 1
+    state.counter += 1
+
+    assert state.chat_history == ["message1"]
+    assert state.items == {"key1": 1}
+    assert state.counter == 1
+
+
+def test_base_state_field_with_default_factory_creates_separate_instances():
+    """Test that Field(default_factory=...) creates separate instances for each state."""
+
+    # GIVEN a state class with Field(default_factory=list)
+    class TestState(BaseState):
+        items: List[str] = Field(default_factory=list)
+
+    # WHEN we create two state instances
+    state1 = TestState()
+    state2 = TestState()
+
+    # THEN they should have separate list instances
+    assert state1.items is not state2.items
+
+    # AND modifying one should not affect the other
+    state1.items.append("item1")
+    assert state1.items == ["item1"]
+    assert state2.items == []
+
+
+class BlockingValue:
+    """A value that blocks during deepcopy until signaled to proceed."""
+
+    def __init__(self, entered_event: threading.Event, proceed_event: threading.Event):
+        self.entered_event = entered_event
+        self.proceed_event = proceed_event
+
+    def __deepcopy__(self, memo: Any) -> "BlockingValue":
+        self.entered_event.set()
+        self.proceed_event.wait(timeout=5.0)
+        return BlockingValue(self.entered_event, self.proceed_event)
+
+
+def test_state_snapshot__concurrent_mutation_during_deepcopy():
+    """Test that concurrent mutations during deepcopy don't cause RuntimeError."""
+
+    # GIVEN a state with a dict containing a blocking value
+    class TestState(BaseState):
+        data: Dict[str, Any] = Field(default_factory=dict)
+
+    state = TestState()
+
+    entered_event = threading.Event()
+    proceed_event = threading.Event()
+    state.data["blocking"] = BlockingValue(entered_event, proceed_event)
+    state.data["other"] = "value"
+
+    snapshot_exception: List[Exception] = []
+    mutation_completed = threading.Event()
+
+    def snapshot_thread_fn() -> None:
+        try:
+            with state.__lock__:
+                deepcopy(state)
+        except Exception as e:
+            snapshot_exception.append(e)
+
+    def mutation_thread_fn() -> None:
+        state.data["new_key"] = "new_value"
+        mutation_completed.set()
+
+    # WHEN we start a snapshot (deepcopy) in one thread
+    snapshot_thread = threading.Thread(target=snapshot_thread_fn)
+    snapshot_thread.start()
+
+    # AND wait for the deepcopy to be in progress (blocked on our blocking value)
+    entered_event.wait(timeout=5.0)
+
+    # AND try to mutate the dict from another thread
+    mutation_thread = threading.Thread(target=mutation_thread_fn)
+    mutation_thread.start()
+
+    # THEN the mutation should block waiting for the lock (not complete immediately)
+    mutation_completed.wait(timeout=0.2)
+    mutation_blocked = not mutation_completed.is_set()
+
+    # AND when we allow the deepcopy to proceed
+    proceed_event.set()
+    snapshot_thread.join(timeout=5.0)
+    mutation_thread.join(timeout=5.0)
+
+    # THEN the mutation should have been blocked by the lock
+    assert mutation_blocked, "Mutation should block while deepcopy holds the lock"
+
+    # AND no exception should have been raised during snapshot
+    assert len(snapshot_exception) == 0, f"Snapshot raised exception: {snapshot_exception}"
+
+
+def test_state_deepcopy__cloned_state_uses_own_snapshot_callback():
+    """Test that deepcopied state's snapshottable containers use the clone's callback."""
+
+    # GIVEN a state with a snapshottable dict attribute
+    original_snapshot_count = 0
+    clone_snapshot_count = 0
+
+    class TestState(BaseState):
+        data: Dict[str, int] = Field(default_factory=dict)
+
+    state = TestState()
+    state.data["key1"] = 1
+
+    def original_callback(state_copy: BaseState, deltas: List[StateDelta]) -> None:
+        nonlocal original_snapshot_count
+        original_snapshot_count += 1
+
+    state.__snapshot_callback__ = original_callback
+
+    # WHEN we deepcopy the state
+    cloned_state = deepcopy(state)
+
+    def clone_callback(state_copy: BaseState, deltas: List[StateDelta]) -> None:
+        nonlocal clone_snapshot_count
+        clone_snapshot_count += 1
+
+    cloned_state.__snapshot_callback__ = clone_callback
+
+    # AND reset counters
+    original_snapshot_count = 0
+    clone_snapshot_count = 0
+
+    # AND mutate the cloned state's snapshottable dict
+    cloned_state.data["key2"] = 2
+
+    # THEN only the clone's callback should be invoked
+    assert clone_snapshot_count == 1, "Clone's callback should be invoked"
+    assert original_snapshot_count == 0, "Original's callback should not be invoked"
+
+
+def test_state_snapshot__top_level_attribute_assignment_blocks_during_deepcopy():
+    """Test that top-level attribute assignments block while deepcopy holds the lock."""
+
+    # GIVEN a state with a blocking value in a dict attribute
+    class TestState(BaseState):
+        data: Dict[str, Any] = Field(default_factory=dict)
+        counter: int = 0
+
+    state = TestState()
+
+    entered_event = threading.Event()
+    proceed_event = threading.Event()
+    state.data["blocking"] = BlockingValue(entered_event, proceed_event)
+
+    mutation_completed = threading.Event()
+
+    def snapshot_thread_fn() -> None:
+        with state.__lock__:
+            deepcopy(state)
+
+    def mutation_thread_fn() -> None:
+        state.__is_quiet__ = True
+        state.counter = 42
+        mutation_completed.set()
+
+    # WHEN we start a snapshot (deepcopy) in one thread
+    snapshot_thread = threading.Thread(target=snapshot_thread_fn)
+    snapshot_thread.start()
+
+    # AND wait for the deepcopy to be in progress (blocked on our blocking value)
+    entered_event.wait(timeout=5.0)
+
+    # AND try to assign a top-level attribute from another thread
+    mutation_thread = threading.Thread(target=mutation_thread_fn)
+    mutation_thread.start()
+
+    # THEN the mutation should block waiting for the lock (not complete immediately)
+    mutation_completed.wait(timeout=0.2)
+    mutation_blocked = not mutation_completed.is_set()
+
+    # AND when we allow the deepcopy to proceed
+    proceed_event.set()
+    snapshot_thread.join(timeout=5.0)
+    mutation_thread.join(timeout=5.0)
+
+    # THEN the mutation should have been blocked by the lock
+    assert mutation_blocked, "Top-level attribute assignment should block while deepcopy holds the lock"
+
+
+def test_base_state_chat_history_with_default_factory_initializes_to_list():
+    """
+    Tests that a chat_history state variable with Optional[list[ChatMessage]] = Field(default_factory=list)
+    initializes to an empty list instead of None.
+    """
+
+    # GIVEN a state class with chat_history using Field(default_factory=list)
+    class TestState(BaseState):
+        chat_history: Optional[List[ChatMessage]] = Field(default_factory=list)  # type: ignore[arg-type]
+
+    # WHEN we create a state instance without providing a value
+    state = TestState()
+
+    # THEN the chat_history should be an empty list, not None
+    assert state.chat_history is not None
+    assert isinstance(state.chat_history, list)
+    assert state.chat_history == []
+
+    # AND we should be able to append ChatMessage objects to it
+    chat_history = state.chat_history
+    chat_history.append(ChatMessage(role="USER", text="Hello"))
+    assert len(chat_history) == 1
+    assert chat_history[0].role == "USER"
+    assert chat_history[0].text == "Hello"
+
+
+def test_base_state_chat_history_with_default_factory_creates_separate_instances():
+    """
+    Tests that Field(default_factory=list) creates separate list instances for each state,
+    avoiding the mutable default argument issue.
+    """
+
+    # GIVEN a state class with chat_history using Field(default_factory=list)
+    class TestState(BaseState):
+        chat_history: Optional[List[ChatMessage]] = Field(default_factory=list)  # type: ignore[arg-type]
+
+    # WHEN we create two state instances
+    state1 = TestState()
+    state2 = TestState()
+
+    # THEN they should have separate list instances
+    assert state1.chat_history is not state2.chat_history
+
+    # AND modifying one should not affect the other
+    chat_history1 = state1.chat_history
+    chat_history2 = state2.chat_history
+    assert chat_history1 is not None
+    assert chat_history2 is not None
+    chat_history1.append(ChatMessage(role="USER", text="Message 1"))
+    chat_history2.append(ChatMessage(role="ASSISTANT", text="Message 2"))
+
+    assert len(chat_history1) == 1
+    assert len(chat_history2) == 1
+    assert chat_history1[0].text == "Message 1"
+    assert chat_history2[0].text == "Message 2"

vellum/workflows/tests/test_dataset_row.py

@@ -152,7 +152,7 @@ def test_dataset_row_with_dict_inputs():
 
 def test_dataset_row_with_node_output_mocks():
     """
-    Test that DatasetRow can be created with node_output_mocks and properly serialized.
+    Test that DatasetRow can be created with mocks and properly serialized.
     """
 
     # GIVEN a node with outputs
@@ -168,7 +168,7 @@ def test_dataset_row_with_node_output_mocks():
 
     test_inputs = TestInputs(message="test message")
 
-    dataset_row = DatasetRow(label="test_with_mocks", inputs=test_inputs, node_output_mocks=[mock_output])
+    dataset_row = DatasetRow(label="test_with_mocks", inputs=test_inputs, mocks=[mock_output])
 
     serialized_dict = dataset_row.model_dump()
 
@@ -176,14 +176,15 @@ def test_dataset_row_with_node_output_mocks():
     assert serialized_dict["label"] == "test_with_mocks"
     assert serialized_dict["inputs"]["message"] == "test message"
 
-    # AND the node_output_mocks should be present in the serialized dict
-    assert "node_output_mocks" in serialized_dict
-    assert serialized_dict["node_output_mocks"] is not None
-    assert len(serialized_dict["node_output_mocks"]) == 1
+    # AND the mocks should be present in the serialized dict
+    assert "mocks" in serialized_dict
+    assert serialized_dict["mocks"] is not None
+    assert len(serialized_dict["mocks"]) == 1
 
     # AND the mock output should be serialized as a dict with the correct structure
-    mock_data = serialized_dict["node_output_mocks"][0]
+    mock_data = serialized_dict["mocks"][0]
     assert mock_data == {
+        "type": "NODE_EXECUTION",
         "node_id": str(DummyNode.__id__),
         "when_condition": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": True}},
         "then_outputs": {"result": "mocked output"},

vellum/workflows/tests/test_sandbox.py

@@ -129,12 +129,15 @@ def test_sandbox_runner_with_workflow_trigger(mock_logger):
         class Outputs(BaseWorkflow.Outputs):
             final_output = StartNode.Outputs.result
 
-    # AND a dataset with workflow_trigger
+    # AND a trigger instance
+    trigger_instance = MySchedule(current_run_at=datetime.min, next_run_at=datetime.now())
+
+    # AND a dataset with workflow_trigger instance
     dataset = [
         DatasetRow(
             label="test_row",
             inputs={"current_run_at": datetime.min, "next_run_at": datetime.now()},
-            workflow_trigger=MySchedule,
+            workflow_trigger=trigger_instance,
         ),
     ]
 
@@ -151,13 +154,99 @@ def test_sandbox_runner_with_workflow_trigger(mock_logger):
         "final_output: 0001-01-01 00:00:00",
     ]
 
-    # AND the dataset row should still have the trigger class
-    assert dataset[0].workflow_trigger == MySchedule
+    # AND the dataset row should have the trigger instance
+    assert dataset[0].workflow_trigger == trigger_instance
+    assert isinstance(dataset[0].workflow_trigger, MySchedule)
+
+
+def test_sandbox_runner_with_trigger_instance(mock_logger):
+    """
+    Test that WorkflowSandboxRunner can run with DatasetRow containing trigger instance.
+    """
+
+    # GIVEN we capture the logs to stdout
+    logs = []
+    mock_logger.return_value.info.side_effect = lambda msg: logs.append(msg)
+
+    # AND a trigger class
+    class MySchedule(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "* * * * *"
+            timezone = "UTC"
+
+    # AND a workflow that uses the trigger
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result = MySchedule.current_run_at
+
+    class Workflow(BaseWorkflow):
+        graph = MySchedule >> StartNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_output = StartNode.Outputs.result
+
+    # AND a trigger instance
+    trigger_instance = MySchedule(current_run_at=datetime.min, next_run_at=datetime.now())
+
+    # AND a dataset with trigger instance
+    dataset = [
+        DatasetRow(
+            label="test_row_with_instance",
+            inputs={"current_run_at": datetime.min, "next_run_at": datetime.now()},
+            workflow_trigger=trigger_instance,
+        ),
+    ]
+
+    # WHEN we run the sandbox with the DatasetRow containing trigger instance
+    runner = WorkflowSandboxRunner(workflow=Workflow(), dataset=dataset)
+    runner.run()
+
+    # THEN the workflow should run successfully
+    assert logs == [
+        "Just started Node: StartNode",
+        "Just finished Node: StartNode",
+        "Workflow fulfilled!",
+        "----------------------------------",
+        "final_output: 0001-01-01 00:00:00",
+    ]
+
+    # AND the dataset row should have the trigger instance
+    assert dataset[0].workflow_trigger == trigger_instance
+    assert isinstance(dataset[0].workflow_trigger, MySchedule)
+
+
+def test_dataset_row_serialization_with_workflow_trigger():
+    """
+    Test that DatasetRow serializes workflow_trigger field to workflow_trigger_id.
+    """
+
+    # GIVEN a trigger class
+    class MySchedule(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "* * * * *"
+            timezone = "UTC"
+
+    # AND a trigger instance
+    trigger_instance = MySchedule(current_run_at=datetime.min, next_run_at=datetime.now())
+
+    # AND a DatasetRow constructed with workflow_trigger
+    dataset_row = DatasetRow(
+        label="test_serialization",
+        inputs={"foo": "bar"},
+        workflow_trigger=trigger_instance,
+    )
+
+    # WHEN we serialize the DatasetRow
+    serialized = dataset_row.model_dump()
+
+    # THEN the serialized dict should contain workflow_trigger_id
+    assert "workflow_trigger_id" in serialized
+    assert serialized["workflow_trigger_id"] == str(MySchedule.__id__)
 
 
 def test_sandbox_runner_with_node_output_mocks(mock_logger, mocker):
     """
-    Tests that WorkflowSandboxRunner passes node_output_mocks from DatasetRow to workflow.stream().
+    Tests that WorkflowSandboxRunner passes mocks from DatasetRow to workflow.stream().
     """
 
     class Inputs(BaseInputs):
@@ -175,12 +264,12 @@ def test_sandbox_runner_with_node_output_mocks(mock_logger, mocker):
 
     mock_outputs = TestNode.Outputs(result="mocked_result")
 
-    # AND a dataset with node_output_mocks
+    # AND a dataset with mocks
     dataset = [
         DatasetRow(
             label="test_with_mocks",
             inputs={"message": "test"},
-            node_output_mocks=[mock_outputs],
+            mocks=[mock_outputs],
         ),
     ]
 
@@ -189,7 +278,7 @@ def test_sandbox_runner_with_node_output_mocks(mock_logger, mocker):
     stream_mock = MagicMock(return_value=original_stream(inputs=Inputs(message="test")))
     mocker.patch.object(workflow_instance, "stream", stream_mock)
 
-    # WHEN we run the sandbox with the DatasetRow containing node_output_mocks
+    # WHEN we run the sandbox with the DatasetRow containing mocks
     runner = WorkflowSandboxRunner(workflow=workflow_instance, dataset=dataset)
     runner.run()
 

vellum/workflows/triggers/__init__.py

@@ -1,6 +1,7 @@
 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
 
-__all__ = ["BaseTrigger", "IntegrationTrigger", "ManualTrigger", "ScheduleTrigger"]
+__all__ = ["BaseTrigger", "ChatMessageTrigger", "IntegrationTrigger", "ManualTrigger", "ScheduleTrigger"]