vellum-ai 0.14.84__tar.gz → 0.14.86__tar.gz

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/PKG-INFO

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

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/pyproject.toml

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

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/client/core/client_wrapper.py

@@ -16,10 +16,10 @@ class BaseClientWrapper:
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
-            "User-Agent": "vellum-ai/0.14.84",
+            "User-Agent": "vellum-ai/0.14.86",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.84",
+            "X-Fern-SDK-Version": "0.14.86",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/client/types/execution_thinking_vellum_value.py

@@ -19,7 +19,7 @@ class ExecutionThinkingVellumValue(UniversalBaseModel):
 
     name: str
     type: typing.Literal["THINKING"] = "THINKING"
-    value: StringVellumValue
+    value: typing.List[StringVellumValue]
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/client/types/thinking_vellum_value.py

@@ -13,7 +13,7 @@ class ThinkingVellumValue(UniversalBaseModel):
     """
 
     type: typing.Literal["THINKING"] = "THINKING"
-    value: StringVellumValue
+    value: typing.List[StringVellumValue]
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/client/types/thinking_vellum_value_request.py

@@ -13,7 +13,7 @@ class ThinkingVellumValueRequest(UniversalBaseModel):
     """
 
     type: typing.Literal["THINKING"] = "THINKING"
-    value: StringVellumValueRequest
+    value: typing.List[StringVellumValueRequest]
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/nodes/core/inline_subworkflow_node/node.py

@@ -84,10 +84,13 @@ class InlineSubworkflowNode(
             )
 
         outputs: Optional[BaseOutputs] = None
+        exception: Optional[NodeException] = None
         fulfilled_output_names: Set[str] = set()
 
         for event in subworkflow_stream:
             self._context._emit_subworkflow_event(event)
+            if exception:
+                continue
 
             if not is_workflow_event(event):
                 continue
@@ -101,7 +104,15 @@ class InlineSubworkflowNode(
             elif event.name == "workflow.execution.fulfilled":
                 outputs = event.outputs
             elif event.name == "workflow.execution.rejected":
-                raise NodeException.of(event.error)
+                exception = NodeException.of(event.error)
+            elif event.name == "workflow.execution.paused":
+                exception = NodeException(
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
+                    message="Subworkflow unexpectedly paused",
+                )
+
+        if exception:
+            raise exception
 
         if outputs is None:
             raise NodeException(

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/nodes/displayable/inline_prompt_node/node.py

@@ -64,8 +64,9 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
             elif output.type == "FUNCTION_CALL":
                 string_outputs.append(output.value.model_dump_json(indent=4))
             elif output.type == "THINKING":
-                if output.value.type == "STRING":
-                    string_outputs.append(output.value.value)
+                for thinking_item in output.value:
+                    if thinking_item.type == "STRING":
+                        string_outputs.append(thinking_item.value)
             else:
                 string_outputs.append(output.value.message)
 

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/nodes/displayable/prompt_deployment_node/node.py

@@ -66,8 +66,9 @@ class PromptDeploymentNode(BasePromptDeploymentNode[StateType]):
             elif output.type == "FUNCTION_CALL":
                 string_outputs.append(output.value.model_dump_json(indent=4))
             elif output.type == "THINKING":
-                if output.value.type == "STRING":
-                    string_outputs.append(output.value.value)
+                for thinking_item in output.value:
+                    if thinking_item.type == "STRING":
+                        string_outputs.append(thinking_item.value)
             else:
                 string_outputs.append(output.value.message)
 

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/nodes/displayable/tool_calling_node/node.py

@@ -1,8 +1,9 @@
-from typing import ClassVar, List, Optional
+from typing import ClassVar, Iterator, List, Optional, Set
 
 from vellum import ChatMessage, PromptBlock
 from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.events.workflow import is_workflow_event
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.graph.graph import Graph
 from vellum.workflows.inputs.base import BaseInputs
@@ -12,10 +13,11 @@ from vellum.workflows.nodes.displayable.tool_calling_node.utils import (
     create_tool_router_node,
     get_function_name,
 )
-from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.core import EntityInputsInterface, Tool
+from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 
 class ToolCallingNode(BaseNode):
@@ -47,12 +49,12 @@ class ToolCallingNode(BaseNode):
         text: str
         chat_history: List[ChatMessage]
 
-    def run(self) -> Outputs:
+    def run(self) -> Iterator[BaseOutput]:
         """
-        Run the tool calling workflow.
+        Run the tool calling workflow with streaming support.
 
         This dynamically builds a graph with router and function nodes,
-        then executes the workflow.
+        then executes the workflow and streams chat_history updates.
         """
 
         self._build_graph()
@@ -77,24 +79,55 @@ class ToolCallingNode(BaseNode):
                 context=WorkflowContext.create_from(self._context),
             )
 
-            terminal_event = subworkflow.run()
-
-            if terminal_event.name == "workflow.execution.paused":
-                raise NodeException(
-                    code=WorkflowErrorCode.INVALID_OUTPUTS,
-                    message="Subworkflow unexpectedly paused",
-                )
-            elif terminal_event.name == "workflow.execution.fulfilled":
-                node_outputs = self.Outputs(
-                    text=terminal_event.outputs.text,
-                    chat_history=terminal_event.outputs.chat_history,
-                )
+            subworkflow_stream = subworkflow.stream(
+                event_filter=all_workflow_event_filter,
+                node_output_mocks=self._context._get_all_node_output_mocks(),
+            )
 
-                return node_outputs
-            elif terminal_event.name == "workflow.execution.rejected":
-                raise NodeException(message=terminal_event.error.message, code=terminal_event.error.code)
+        outputs: Optional[BaseOutputs] = None
+        exception: Optional[NodeException] = None
+        fulfilled_output_names: Set[str] = set()
+
+        # Yield initiated event for chat_history
+        yield BaseOutput(name="chat_history")
+
+        for event in subworkflow_stream:
+            self._context._emit_subworkflow_event(event)
+
+            if not is_workflow_event(event):
+                continue
+            if event.workflow_definition != ToolCallingWorkflow:
+                continue
+
+            if event.name == "workflow.execution.streaming":
+                if event.output.name == "chat_history":
+                    if event.output.is_fulfilled:
+                        fulfilled_output_names.add(event.output.name)
+                    yield event.output
+                elif event.output.name == "text":
+                    if event.output.is_fulfilled:
+                        fulfilled_output_names.add(event.output.name)
+                    yield event.output
+            elif event.name == "workflow.execution.fulfilled":
+                outputs = event.outputs
+            elif event.name == "workflow.execution.rejected":
+                exception = NodeException.of(event.error)
+
+        if exception:
+            raise exception
+
+        if outputs is None:
+            raise NodeException(
+                message="Expected to receive outputs from Tool Calling Workflow",
+                code=WorkflowErrorCode.INVALID_OUTPUTS,
+            )
 
-            raise Exception(f"Unexpected workflow event: {terminal_event.name}")
+        for output_descriptor, output_value in outputs:
+            if output_descriptor.name not in fulfilled_output_names:
+                yield BaseOutput(
+                    name=output_descriptor.name,
+                    value=output_value,
+                )
 
     def _build_graph(self) -> None:
         self.tool_router_node = create_tool_router_node(

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/nodes/displayable/tool_calling_node/tests/test_node.py

@@ -114,7 +114,7 @@ def test_tool_calling_node_inline_workflow_context():
     function_node.state = TestState(meta=StateMeta(node_outputs={tool_router_node.Outputs.text: '{"arguments": {}}'}))
 
     # WHEN the function node runs
-    outputs = function_node.run()
+    outputs = list(function_node.run())
 
     # THEN the workflow should have run successfully
     assert outputs is not None

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/runner/runner.py

@@ -15,6 +15,7 @@ from typing import (
     Generic,
     Iterable,
     Iterator,
+    List,
     Optional,
     Sequence,
     Set,
@@ -67,7 +68,9 @@ from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references import ExternalInputReference, OutputReference
+from vellum.workflows.references.state_value import StateValueReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.delta import StateDelta
 from vellum.workflows.types.generics import InputsType, OutputsType, StateType
 
 if TYPE_CHECKING:
@@ -175,12 +178,16 @@ class WorkflowRunner(Generic[StateType]):
         setattr(
             self._initial_state,
             "__snapshot_callback__",
-            lambda s: self._snapshot_state(s),
+            lambda s, d: self._snapshot_state(s, d),
         )
         self.workflow.context._register_event_queue(self._workflow_event_inner_queue)
         self.workflow.context._register_node_output_mocks(node_output_mocks or [])
 
-    def _snapshot_state(self, state: StateType) -> StateType:
+        self._outputs_listening_to_state = [
+            descriptor for descriptor in self.workflow.Outputs if isinstance(descriptor.instance, StateValueReference)
+        ]
+
+    def _snapshot_state(self, state: StateType, deltas: List[StateDelta]) -> StateType:
         self._workflow_event_inner_queue.put(
             WorkflowExecutionSnapshottedEvent(
                 trace_id=self._execution_context.trace_id,
@@ -192,6 +199,31 @@ class WorkflowRunner(Generic[StateType]):
                 parent=self._execution_context.parent_context,
             )
         )
+
+        delta_names = [d.name for d in deltas]
+        for descriptor in self._outputs_listening_to_state:
+            if not isinstance(descriptor.instance, StateValueReference):
+                continue
+
+            if descriptor.instance.name not in delta_names:
+                continue
+
+            resolved_delta = descriptor.instance.resolve(state)
+            if resolved_delta is undefined:
+                continue
+
+            self._workflow_event_outer_queue.put(
+                self._stream_workflow_event(
+                    BaseOutput(
+                        name=descriptor.name,
+                        # We may want to switch this to using the delta from the state delta
+                        # instead of the state value. In the short term, we need to show
+                        # the full conversation of the tool calling node's delta's.
+                        delta=resolved_delta,
+                    )
+                )
+            )
+
         self.workflow._store.append_state_snapshot(state)
         self._background_thread_queue.put(state)
         return state

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/state/base.py

@@ -7,7 +7,21 @@ import logging
 from queue import Queue
 from threading import Lock
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    SupportsIndex,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, field_serializer, field_validator
@@ -18,6 +32,7 @@ from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.constants import undefined
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
+from vellum.workflows.state.delta import AppendStateDelta, SetStateDelta, StateDelta
 from vellum.workflows.types.definition import CodeResourceDefinition, serialize_type_encoder_with_id
 from vellum.workflows.types.generics import StateType, import_workflow_class, is_workflow_class
 from vellum.workflows.types.stack import Stack
@@ -31,16 +46,12 @@ logger = logging.getLogger(__name__)
 
 
 class _Snapshottable:
-    _snapshot_callback: Callable[[], None]
+    _snapshot_callback: Callable[[Optional[StateDelta]], None]
+    _path: str
 
-    def __deepcopy__(self, memo: Any) -> "_Snapshottable":
-        return deepcopy_with_exclusions(
-            self,
-            memo=memo,
-            exclusions={
-                "_snapshot_callback": self._snapshot_callback,
-            },
-        )
+    def __bind__(self, path: str, snapshot_callback: Callable[[Optional[StateDelta]], None]) -> None:
+        self._snapshot_callback = snapshot_callback
+        self._path = path
 
 
 @dataclass_transform(kw_only_default=True)
@@ -71,10 +82,44 @@ class _BaseStateMeta(type):
 class _SnapshottableDict(dict, _Snapshottable):
     def __setitem__(self, key: Any, value: Any) -> None:
         super().__setitem__(key, value)
-        self._snapshot_callback()
+        self._snapshot_callback(SetStateDelta(name=f"{self._path}.{key}", delta=value))
+
+    def __deepcopy__(self, memo: Any) -> "_SnapshottableDict":
+        y: dict = {}
+        memo[id(self)] = y
+        for key, value in self.items():
+            y[deepcopy(key, memo)] = deepcopy(value, memo)
+
+        y = _SnapshottableDict(y)
+        y.__bind__(self._path, self._snapshot_callback)
+        memo[id(self)] = y
+        return y
+
+
+class _SnapshottableList(list, _Snapshottable):
+    def __setitem__(self, index: Union[SupportsIndex, slice], value: Any) -> None:
+        super().__setitem__(index, value)
+        if isinstance(index, int):
+            self._snapshot_callback(SetStateDelta(name=f"{self._path}.{index}", delta=value))
+
+    def append(self, value: Any) -> None:
+        super().append(value)
+        self._snapshot_callback(AppendStateDelta(name=self._path, delta=value))
+
+    def __deepcopy__(self, memo: Any) -> "_SnapshottableList":
+        y: list = []
+        memo[id(self)] = y
+        append = y.append
+        for a in self:
+            append(deepcopy(a, memo))
 
+        y = _SnapshottableList(y)
+        y.__bind__(self._path, self._snapshot_callback)
+        memo[id(self)] = y
+        return y
 
-def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> Any:
+
+def _make_snapshottable(path: str, value: Any, snapshot_callback: Callable[[Optional[StateDelta]], None]) -> Any:
     """
     Edits any value to make it snapshottable on edit. Made as a separate function from `BaseState` to
     avoid namespace conflicts with subclasses.
@@ -84,9 +129,14 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
 
     if isinstance(value, dict):
         snapshottable_dict = _SnapshottableDict(value)
-        snapshottable_dict._snapshot_callback = snapshot_callback
+        snapshottable_dict.__bind__(path, snapshot_callback)
         return snapshottable_dict
 
+    if isinstance(value, list):
+        snapshottable_list = _SnapshottableList(value)
+        snapshottable_list.__bind__(path, snapshot_callback)
+        return snapshottable_list
+
     return value
 
 
@@ -253,14 +303,14 @@ class StateMeta(UniversalBaseModel):
     node_outputs: Dict[OutputReference, Any] = field(default_factory=dict)
     node_execution_cache: NodeExecutionCache = field(default_factory=NodeExecutionCache)
     parent: Optional["BaseState"] = None
-    __snapshot_callback__: Optional[Callable[[], None]] = field(init=False, default=None)
+    __snapshot_callback__: Optional[Callable[[Optional[StateDelta]], None]] = field(init=False, default=None)
 
     def model_post_init(self, context: Any) -> None:
         self.__snapshot_callback__ = None
 
-    def add_snapshot_callback(self, callback: Callable[[], None]) -> None:
-        self.node_outputs = _make_snapshottable(self.node_outputs, callback)
-        self.external_inputs = _make_snapshottable(self.external_inputs, callback)
+    def add_snapshot_callback(self, callback: Callable[[Optional[StateDelta]], None]) -> None:
+        self.node_outputs = _make_snapshottable("meta.node_outputs", self.node_outputs, callback)
+        self.external_inputs = _make_snapshottable("meta.external_inputs", self.external_inputs, callback)
         self.__snapshot_callback__ = callback
 
     def __setattr__(self, name: str, value: Any) -> None:
@@ -270,7 +320,7 @@ class StateMeta(UniversalBaseModel):
 
         super().__setattr__(name, value)
         if callable(self.__snapshot_callback__):
-            self.__snapshot_callback__()
+            self.__snapshot_callback__(SetStateDelta(name=f"meta.{name}", delta=value))
 
     @field_serializer("workflow_definition")
     def serialize_workflow_definition(self, workflow_definition: Type["BaseWorkflow"], _info: Any) -> Dict[str, Any]:
@@ -428,11 +478,15 @@ class BaseState(metaclass=_BaseStateMeta):
 
     __lock__: Lock = field(init=False)
     __is_quiet__: bool = field(init=False)
-    __snapshot_callback__: Callable[["BaseState"], None] = field(init=False)
+    __is_atomic__: bool = field(init=False)
+    __snapshot_callback__: Callable[["BaseState", List[StateDelta]], None] = field(init=False)
+    __deltas__: List[StateDelta] = field(init=False)
 
     def __init__(self, meta: Optional[StateMeta] = None, **kwargs: Any) -> None:
         self.__is_quiet__ = True
-        self.__snapshot_callback__ = lambda state: None
+        self.__is_atomic__ = False
+        self.__snapshot_callback__ = lambda state, deltas: None
+        self.__deltas__ = []
         self.__lock__ = Lock()
 
         self.meta = meta or StateMeta()
@@ -442,7 +496,7 @@ class BaseState(metaclass=_BaseStateMeta):
         for name, value in self.__class__.__dict__.items():
             if not name.startswith("_") and name != "meta":
                 # Bypass __is_quiet__ instead of `setattr`
-                snapshottable_value = _make_snapshottable(value, self.__snapshot__)
+                snapshottable_value = _make_snapshottable(name, value, self.__snapshot__)
                 super().__setattr__(name, snapshottable_value)
 
         for name, value in kwargs.items():
@@ -494,15 +548,15 @@ class BaseState(metaclass=_BaseStateMeta):
     def __getitem__(self, key: str) -> Any:
         return self.__dict__[key]
 
-    def __setattr__(self, name: str, value: Any) -> None:
+    def __setattr__(self, name: str, delta: Any) -> None:
         if name.startswith("_"):
-            super().__setattr__(name, value)
+            super().__setattr__(name, delta)
             return
 
-        snapshottable_value = _make_snapshottable(value, self.__snapshot__)
+        snapshottable_value = _make_snapshottable(name, delta, self.__snapshot__)
         super().__setattr__(name, snapshottable_value)
         self.meta.updated_ts = datetime_now()
-        self.__snapshot__()
+        self.__snapshot__(SetStateDelta(name=name, delta=delta))
 
     def __add__(self, other: StateType) -> StateType:
         """
@@ -531,7 +585,7 @@ class BaseState(metaclass=_BaseStateMeta):
 
         return cast(StateType, latest_state)
 
-    def __snapshot__(self) -> None:
+    def __snapshot__(self, delta: Optional[StateDelta] = None) -> None:
         """
         Snapshots the current state to the workflow emitter. The invoked callback is overridden by the
         workflow runner.
@@ -539,11 +593,19 @@ class BaseState(metaclass=_BaseStateMeta):
         if self.__is_quiet__:
             return
 
+        if delta:
+            self.__deltas__.append(delta)
+
+        if self.__is_atomic__:
+            return
+
         try:
-            self.__snapshot_callback__(deepcopy(self))
+            self.__snapshot_callback__(deepcopy(self), self.__deltas__)
         except Exception:
             logger.exception("Failed to snapshot Workflow state.")
 
+        self.__deltas__.clear()
+
     @contextmanager
     def __quiet__(self):
         prev = self.__is_quiet__
@@ -555,12 +617,12 @@ class BaseState(metaclass=_BaseStateMeta):
 
     @contextmanager
     def __atomic__(self):
-        prev = self.__is_quiet__
-        self.__is_quiet__ = True
+        prev = self.__is_atomic__
+        self.__is_atomic__ = True
         try:
             yield
         finally:
-            self.__is_quiet__ = prev
+            self.__is_atomic__ = prev
             self.__snapshot__()
 
     @classmethod

vellum_ai-0.14.86/src/vellum/workflows/state/delta.py

@@ -0,0 +1,20 @@
+from typing import Any, Literal, Union
+
+from vellum.core.pydantic_utilities import UniversalBaseModel
+
+
+class BaseStateDelta(UniversalBaseModel):
+    name: str
+
+
+class SetStateDelta(BaseStateDelta):
+    delta: Any
+    delta_type: Literal["SET_STATE"] = "SET_STATE"
+
+
+class AppendStateDelta(BaseStateDelta):
+    delta: Any
+    delta_type: Literal["APPEND_STATE"] = "APPEND_STATE"
+
+
+StateDelta = Union[SetStateDelta, AppendStateDelta]

{vellum_ai-0.14.84 → vellum_ai-0.14.86}/src/vellum/workflows/state/tests/test_state.py

@@ -2,11 +2,12 @@ import pytest
 from copy import deepcopy
 import json
 from queue import Queue
-from typing import Dict, cast
+from typing import Dict, List, cast
 
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.delta import SetStateDelta, StateDelta
 from vellum.workflows.state.encoder import DefaultStateEncoder
 
 
@@ -28,7 +29,7 @@ class MockState(BaseState):
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.__snapshot_callback__ = lambda _: self.__mock_snapshot__()
+        self.__snapshot_callback__ = lambda _, __: self.__mock_snapshot__()
 
     def __mock_snapshot__(self) -> None:
         self.__snapshot_count__ += 1
@@ -201,7 +202,7 @@ def test_state_snapshot__deepcopy_fails__logs_error(mock_deepcopy, mock_logger):
     # AND we have a snapshot callback that we are tracking
     snapshot_count = 0
 
-    def __snapshot_callback__(state: "BaseState") -> None:
+    def __snapshot_callback__(state: "BaseState", deltas: List[StateDelta]) -> None:
         nonlocal snapshot_count
         snapshot_count += 1
 
@@ -220,8 +221,8 @@ def test_state_snapshot__deepcopy_fails__logs_error(mock_deepcopy, mock_logger):
     mock_deepcopy.side_effect = side_effect
 
     # WHEN we snapshot the state twice
-    state.__snapshot__()
-    state.__snapshot__()
+    state.__snapshot__(SetStateDelta(name="foo", delta="bar"))
+    state.__snapshot__(SetStateDelta(name="foo", delta="baz"))
 
     # THEN we were able to invoke the callback once
     assert snapshot_count == 1