vellum-ai 1.2.5__py3-none-any.whl → 1.3.1__py3-none-any.whl

vellum/workflows/resolvers/tests/test_resolver.py

@@ -2,10 +2,18 @@ from datetime import datetime
 from unittest.mock import Mock
 from uuid import uuid4
 
+from vellum.client.types.span_link import SpanLink
+from vellum.client.types.vellum_code_resource_definition import VellumCodeResourceDefinition
 from vellum.client.types.workflow_execution_detail import WorkflowExecutionDetail
+from vellum.client.types.workflow_execution_initiated_body import WorkflowExecutionInitiatedBody
+from vellum.client.types.workflow_execution_initiated_event import WorkflowExecutionInitiatedEvent
+from vellum.client.types.workflow_execution_span import WorkflowExecutionSpan
+from vellum.client.types.workflow_execution_span_attributes import WorkflowExecutionSpanAttributes
+from vellum.client.types.workflow_parent_context import WorkflowParentContext
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.resolvers.resolver import VellumResolver
+from vellum.workflows.resolvers.types import LoadStateResult
 from vellum.workflows.state.base import BaseState, NodeExecutionCache
 from vellum.workflows.state.context import WorkflowContext
 
@@ -14,6 +22,7 @@ def test_load_state_with_context_success():
     """Test load_state successfully loads state when context and client are available."""
     resolver = VellumResolver()
     execution_id = uuid4()
+    root_execution_id = uuid4()
 
     class TestState(BaseState):
         test_key: str = "test_value"
@@ -22,12 +31,14 @@ def test_load_state_with_context_success():
         pass
 
     # GIVEN a state dictionary that matches what the resolver expects
+    prev_id = str(uuid4())
+    prev_span_id = str(uuid4())
     state_dict = {
         "test_key": "test_value",
         "meta": {
             "workflow_definition": "MockWorkflow",
-            "id": str(uuid4()),
-            "span_id": str(uuid4()),
+            "id": prev_id,
+            "span_id": prev_span_id,
             "updated_ts": datetime.now().isoformat(),
             "workflow_inputs": BaseInputs(),
             "external_inputs": {},
@@ -37,22 +48,83 @@ def test_load_state_with_context_success():
         },
     }
 
+    mock_workflow_definition = VellumCodeResourceDefinition(
+        name="TestWorkflow", module=["test", "module"], id=str(uuid4())
+    )
+
+    mock_body = WorkflowExecutionInitiatedBody(workflow_definition=mock_workflow_definition, inputs={})
+
+    previous_trace_id = str(uuid4())
+    root_trace_id = str(uuid4())
+
+    previous_invocation = WorkflowExecutionInitiatedEvent(
+        id=str(uuid4()),
+        timestamp=datetime.now(),
+        trace_id=previous_trace_id,
+        span_id=str(execution_id),
+        body=mock_body,
+        links=[
+            SpanLink(
+                trace_id=previous_trace_id,
+                type="PREVIOUS_SPAN",
+                span_context=WorkflowParentContext(workflow_definition=mock_workflow_definition, span_id=str(uuid4())),
+            ),
+            SpanLink(
+                trace_id=root_trace_id,
+                type="ROOT_SPAN",
+                span_context=WorkflowParentContext(
+                    workflow_definition=mock_workflow_definition, span_id=str(root_execution_id)
+                ),
+            ),
+        ],
+    )
+
+    root_invocation = WorkflowExecutionInitiatedEvent(
+        id=str(uuid4()),
+        timestamp=datetime.now(),
+        trace_id=root_trace_id,
+        span_id=str(root_execution_id),
+        body=mock_body,
+        links=None,  # Root invocation has no links
+    )
+
+    mock_span = WorkflowExecutionSpan(
+        span_id=str(execution_id),  # Use the actual execution_id
+        start_ts=datetime.now(),
+        end_ts=datetime.now(),
+        attributes=WorkflowExecutionSpanAttributes(label="Test Workflow", workflow_id=str(uuid4())),
+        events=[previous_invocation, root_invocation],
+    )
+
     mock_response = WorkflowExecutionDetail(
-        span_id="test-span-id", start=datetime.now(), inputs=[], outputs=[], spans=[], state=state_dict
+        span_id="test-span-id", start=datetime.now(), inputs=[], outputs=[], spans=[mock_span], state=state_dict
     )
 
     mock_client = Mock()
     mock_client.workflow_executions.retrieve_workflow_execution_detail.return_value = mock_response
 
-    # AND context with the test workflow class
+    # AND context with the test workflow class is set up
     context = WorkflowContext(vellum_client=mock_client)
     TestWorkflow(context=context, resolvers=[resolver])
 
+    # WHEN load_state is called
     result = resolver.load_state(previous_execution_id=execution_id)
 
-    # THEN should return an instance of TestWorkflow.State, not BaseState
-    assert isinstance(result, TestState)
-    assert result.test_key == "test_value"
+    # THEN should return LoadStateResult with state and span link info
+    assert isinstance(result, LoadStateResult)
+    assert result.state is not None
+    assert isinstance(result.state, TestState)
+    assert result.state.test_key == "test_value"
+
+    # AND the new state should have different meta IDs than those provided in the loaded state_dict
+    assert str(result.state.meta.id) != prev_id
+    assert str(result.state.meta.span_id) != prev_span_id
+
+    # AND should have span link info
+    assert result.previous_trace_id == previous_invocation.trace_id
+    assert result.previous_span_id == previous_invocation.span_id
+    assert result.root_trace_id == root_invocation.trace_id
+    assert result.root_span_id == root_invocation.span_id
 
     mock_client.workflow_executions.retrieve_workflow_execution_detail.assert_called_once_with(
         execution_id=str(execution_id)

vellum/workflows/resolvers/types.py

@@ -0,0 +1,11 @@
+from typing import NamedTuple
+
+from vellum.workflows.state.base import BaseState
+
+
+class LoadStateResult(NamedTuple):
+    state: BaseState
+    previous_trace_id: str
+    previous_span_id: str
+    root_trace_id: str
+    root_span_id: str

vellum/workflows/runner/runner.py

@@ -46,7 +46,7 @@ from vellum.workflows.events.node import (
     NodeExecutionRejectedBody,
     NodeExecutionStreamingBody,
 )
-from vellum.workflows.events.types import BaseEvent, NodeParentContext, ParentContext, WorkflowParentContext
+from vellum.workflows.events.types import BaseEvent, NodeParentContext, ParentContext, SpanLink, WorkflowParentContext
 from vellum.workflows.events.workflow import (
     WorkflowEventStream,
     WorkflowExecutionFulfilledBody,
@@ -99,6 +99,7 @@ class WorkflowRunner(Generic[StateType]):
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
+        previous_execution_id: Optional[Union[str, UUID]] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
         node_output_mocks: Optional[MockNodeExecutionArg] = None,
         max_concurrency: Optional[int] = None,
@@ -110,6 +111,7 @@ class WorkflowRunner(Generic[StateType]):
         self.workflow = workflow
         self._is_resuming = False
         self._should_emit_initial_state = True
+        self._span_link_info: Optional[Tuple[str, str, str, str]] = None
         if entrypoint_nodes:
             if len(list(entrypoint_nodes)) > 1:
                 raise ValueError("Cannot resume from multiple nodes")
@@ -140,6 +142,28 @@ class WorkflowRunner(Generic[StateType]):
                 if issubclass(ei.inputs_class.__parent_class__, BaseNode)
             ]
             self._is_resuming = True
+        elif previous_execution_id:
+            for resolver in self.workflow.resolvers:
+                try:
+                    load_state_result = resolver.load_state(previous_execution_id)
+                    if load_state_result is not None:
+                        state_class = self.workflow.get_state_class()
+                        if isinstance(load_state_result.state, state_class):
+                            self._initial_state = load_state_result.state
+                            normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
+                            self._initial_state.meta.workflow_inputs = normalized_inputs
+                            self._initial_state.meta.workflow_definition = self.workflow.__class__
+                            self._span_link_info = (
+                                load_state_result.previous_trace_id,
+                                load_state_result.previous_span_id,
+                                load_state_result.root_trace_id,
+                                load_state_result.root_span_id,
+                            )
+                            break
+                except Exception as e:
+                    logger.warning(f"Failed to load state from resolver {type(resolver).__name__}: {e}")
+                    continue
+            self._entrypoints = self.workflow.get_entrypoints()
         else:
             normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
             if state:
@@ -627,6 +651,29 @@ class WorkflowRunner(Generic[StateType]):
         return None
 
     def _initiate_workflow_event(self) -> WorkflowExecutionInitiatedEvent:
+        links: Optional[List[SpanLink]] = None
+
+        if self._span_link_info:
+            previous_trace_id, previous_span_id, root_trace_id, root_span_id = self._span_link_info
+            links = [
+                SpanLink(
+                    trace_id=previous_trace_id,
+                    type="PREVIOUS_SPAN",
+                    span_context=WorkflowParentContext(
+                        workflow_definition=self.workflow.__class__,
+                        span_id=previous_span_id,
+                    ),
+                ),
+                SpanLink(
+                    trace_id=root_trace_id,
+                    type="ROOT_SPAN",
+                    span_context=WorkflowParentContext(
+                        workflow_definition=self.workflow.__class__,
+                        span_id=root_span_id,
+                    ),
+                ),
+            ]
+
         return WorkflowExecutionInitiatedEvent(
             trace_id=self._execution_context.trace_id,
             span_id=self._initial_state.meta.span_id,
@@ -636,6 +683,7 @@ class WorkflowRunner(Generic[StateType]):
                 initial_state=deepcopy(self._initial_state) if self._should_emit_initial_state else None,
             ),
             parent=self._execution_context.parent_context,
+            links=links,
         )
 
     def _stream_workflow_event(self, output: BaseOutput) -> WorkflowExecutionStreamingEvent:

vellum/workflows/state/context.py

@@ -1,15 +1,16 @@
 from functools import cached_property
 from queue import Queue
-from uuid import uuid4
+from uuid import UUID, uuid4
 from typing import TYPE_CHECKING, Dict, List, Optional, Type
 
-from vellum import Vellum
+from vellum import Vellum, __version__
 from vellum.workflows.context import ExecutionContext, get_execution_context, set_execution_context
 from vellum.workflows.events.types import ExternalParentContext
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.utils.uuids import generate_workflow_deployment_prefix
+from vellum.workflows.utils.zip import extract_zip_files
 from vellum.workflows.vellum_client import create_vellum_client
 
 if TYPE_CHECKING:
@@ -40,6 +41,9 @@ class WorkflowContext:
 
         if self._execution_context.parent_context is None:
             self._execution_context.parent_context = ExternalParentContext(span_id=uuid4())
+            # Only generate a new trace_id if one wasn't explicitly provided (i.e., if it's the default zero UUID)
+            if self._execution_context.trace_id == UUID("00000000-0000-0000-0000-000000000000"):
+                self._execution_context.trace_id = uuid4()
             # Propagate the updated context back to the global execution context
             set_execution_context(self._execution_context)
 
@@ -156,23 +160,53 @@ class WorkflowContext:
         Returns:
             BaseWorkflow instance if found, None otherwise
         """
-        if not self.generated_files or not self.namespace:
+        if not self._generated_files or not self._namespace:
             return None
 
         expected_prefix = generate_workflow_deployment_prefix(deployment_name, release_tag)
 
-        workflow_file_key = f"{expected_prefix}/workflow.py"
-        if workflow_file_key not in self.generated_files:
-            return None
+        try:
+            from vellum.workflows.workflows.base import BaseWorkflow
+
+            WorkflowClass = BaseWorkflow.load_from_module(f"{self.namespace}.{expected_prefix}")
+            WorkflowClass.is_dynamic = True
+            workflow_instance = WorkflowClass(context=WorkflowContext.create_from(self), parent_state=state)
+            return workflow_instance
+        except Exception:
+            pass
 
         try:
+            major_version = __version__.split(".")[0]
+            version_range = f">={major_version}.0.0,<={__version__}"
+
+            response = self.vellum_client.workflows.pull(
+                deployment_name,
+                release_tag=release_tag,
+                version=version_range,
+                request_options={"additional_headers": {"X-Vellum-Always-Success": "true"}},
+            )
+
+            if isinstance(response, dict) and response.get("success") is False:
+                return None
+
+            zip_bytes = b"".join(response)
+            pulled_files = extract_zip_files(zip_bytes)
+
+            for file_name, content in pulled_files.items():
+                prefixed_file_name = f"{expected_prefix}/{file_name}"
+                self._generated_files[prefixed_file_name] = content
+
             from vellum.workflows.workflows.base import BaseWorkflow
 
             WorkflowClass = BaseWorkflow.load_from_module(f"{self.namespace}.{expected_prefix}")
+            WorkflowClass.is_dynamic = True
             workflow_instance = WorkflowClass(context=WorkflowContext.create_from(self), parent_state=state)
             return workflow_instance
+
         except Exception:
-            return None
+            pass
+
+        return None
 
     @classmethod
     def create_from(cls, context):

vellum/workflows/utils/zip.py

@@ -0,0 +1,46 @@
+import io
+import zipfile
+
+
+def zip_file_map(file_map: dict[str, str]) -> bytes:
+    """
+    Create a zip file from a dictionary of file names to content.
+
+    Args:
+        file_map: Dictionary mapping file names to their content
+
+    Returns:
+        Bytes representing the zip file
+    """
+    zip_buffer = io.BytesIO()
+
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+        for filename, content in file_map.items():
+            zip_file.writestr(filename, content)
+
+    zip_bytes = zip_buffer.getvalue()
+    zip_buffer.close()
+
+    return zip_bytes
+
+
+def extract_zip_files(zip_bytes: bytes) -> dict[str, str]:
+    """
+    Extract files from a zip archive.
+
+    Args:
+        zip_bytes: Bytes representing the zip file
+
+    Returns:
+        Dictionary mapping file names to their content
+    """
+    zip_buffer = io.BytesIO(zip_bytes)
+    extracted_files = {}
+
+    with zipfile.ZipFile(zip_buffer) as zip_file:
+        for file_name in zip_file.namelist():
+            with zip_file.open(file_name) as source:
+                content = source.read().decode("utf-8")
+                extracted_files[file_name] = content
+
+    return extracted_files

vellum/workflows/workflows/base.py

@@ -364,6 +364,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
+        previous_execution_id: Optional[Union[str, UUID]] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
         node_output_mocks: Optional[MockNodeExecutionArg] = None,
         max_concurrency: Optional[int] = None,
@@ -389,6 +390,9 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         external_inputs: Optional[ExternalInputsArg] = None
             External inputs to pass to the Workflow. Useful for providing human-in-the-loop behavior to the Workflow.
 
+        previous_execution_id: Optional[Union[str, UUID]] = None
+            The execution ID of the previous execution to resume from.
+
         cancel_signal: Optional[ThreadingEvent] = None
             A threading event that can be used to cancel the Workflow Execution.
 
@@ -408,6 +412,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             state=state,
             entrypoint_nodes=entrypoint_nodes,
             external_inputs=external_inputs,
+            previous_execution_id=previous_execution_id,
             cancel_signal=cancel_signal,
             node_output_mocks=node_output_mocks,
             max_concurrency=max_concurrency,
@@ -476,6 +481,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
+        previous_execution_id: Optional[Union[str, UUID]] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
         node_output_mocks: Optional[MockNodeExecutionArg] = None,
         max_concurrency: Optional[int] = None,
@@ -502,6 +508,9 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         external_inputs: Optional[ExternalInputsArg] = None
             External inputs to pass to the Workflow. Useful for providing human-in-the-loop behavior to the Workflow.
 
+        previous_execution_id: Optional[Union[str, UUID]] = None
+            The execution ID of the previous execution to resume from.
+
         cancel_signal: Optional[ThreadingEvent] = None
             A threading event that can be used to cancel the Workflow Execution.
 
@@ -522,6 +531,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             state=state,
             entrypoint_nodes=entrypoint_nodes,
             external_inputs=external_inputs,
+            previous_execution_id=previous_execution_id,
             cancel_signal=cancel_signal,
             node_output_mocks=node_output_mocks,
             max_concurrency=max_concurrency,

{vellum_ai-1.2.5.dist-info → vellum_ai-1.3.1.dist-info}/METADATA

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