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

vellum/workflows/references/trigger.py

@@ -0,0 +1,77 @@
+"""Descriptor for referring to trigger attributes in workflow graphs."""
+
+from __future__ import annotations
+
+from uuid import UUID
+from typing import TYPE_CHECKING, Any, Generic, Optional, Tuple, Type, TypeVar, cast
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.utils.uuids import get_trigger_attribute_id
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+    from vellum.workflows.triggers.base import BaseTrigger
+
+_T = TypeVar("_T")
+
+
+class TriggerAttributeReference(BaseDescriptor[_T], Generic[_T]):
+    """Reference to a trigger attribute defined via type annotations."""
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        types: Tuple[Type[_T], ...],
+        instance: Optional[_T],
+        trigger_class: Type[BaseTrigger],
+    ) -> None:
+        super().__init__(name=name, types=types, instance=instance)
+        self._trigger_class = trigger_class
+
+    @property
+    def trigger_class(self) -> Type[BaseTrigger]:
+        return self._trigger_class
+
+    @property
+    def id(self) -> UUID:
+        """Generate deterministic UUID from trigger class qualname and attribute name."""
+        return get_trigger_attribute_id(self._trigger_class, self.name)
+
+    def resolve(self, state: BaseState) -> _T:
+        trigger_attributes = getattr(state.meta, "trigger_attributes", {})
+        if self in trigger_attributes:
+            return cast(_T, trigger_attributes[self])
+
+        if state.meta.parent:
+            return self.resolve(state.meta.parent)
+
+        if type(None) in self.types:
+            return cast(_T, None)
+
+        raise NodeException(
+            message=f"Missing trigger attribute '{self.name}' for {self._trigger_class.__name__}",
+            code=WorkflowErrorCode.INVALID_INPUTS,
+        )
+
+    def __repr__(self) -> str:
+        return f"{self._trigger_class.__qualname__}.{self.name}"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TriggerAttributeReference):
+            return False
+        return super().__eq__(other) and self._trigger_class == other._trigger_class
+
+    def __hash__(self) -> int:
+        return hash((self._trigger_class, self._name))
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.is_instance_schema(cls)

vellum/workflows/references/workflow_input.py

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Generic, Optional, Tuple, Type, TypeVar, cast
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.types.generics import import_workflow_class
 
 if TYPE_CHECKING:
     from vellum.workflows.inputs.base import BaseInputs
@@ -29,7 +30,10 @@ class WorkflowInputReference(BaseDescriptor[_InputType], Generic[_InputType]):
         return self._inputs_class
 
     def resolve(self, state: "BaseState") -> _InputType:
-        if hasattr(state.meta.workflow_inputs, self._name):
+        if hasattr(state.meta.workflow_inputs, self._name) and (
+            state.meta.workflow_definition == self._inputs_class.__parent_class__
+            or not issubclass(self._inputs_class.__parent_class__, import_workflow_class())
+        ):
             return cast(_InputType, getattr(state.meta.workflow_inputs, self._name))
 
         if state.meta.parent:

vellum/workflows/resolvers/base.py

@@ -1,11 +1,25 @@
 from abc import ABC, abstractmethod
-from typing import Iterator
+from uuid import UUID
+from typing import TYPE_CHECKING, Iterator, Optional, Type, Union
 
 from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.resolvers.types import LoadStateResult
 from vellum.workflows.state.base import BaseState
 
+if TYPE_CHECKING:
+    from vellum.workflows.state.context import WorkflowContext
+    from vellum.workflows.workflows.base import BaseWorkflow
+
 
 class BaseWorkflowResolver(ABC):
+    def __init__(self):
+        self._context: Optional["WorkflowContext"] = None
+        self._workflow_class: Optional[Type["BaseWorkflow"]] = None
+
+    def register_workflow_instance(self, workflow_instance: "BaseWorkflow") -> None:
+        self._workflow_class = type(workflow_instance)
+        self._context = workflow_instance.context
+
     @abstractmethod
     def get_latest_execution_events(self) -> Iterator[WorkflowEvent]:
         pass
@@ -13,3 +27,7 @@ class BaseWorkflowResolver(ABC):
     @abstractmethod
     def get_state_snapshot_history(self) -> Iterator[BaseState]:
         pass
+
+    @abstractmethod
+    def load_state(self, previous_execution_id: Optional[Union[UUID, str]] = None) -> Optional[LoadStateResult]:
+        pass

vellum/workflows/resolvers/resolver.py

@@ -0,0 +1,71 @@
+import logging
+from uuid import UUID
+from typing import Iterator, Optional, Type, Union
+
+from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.nodes.utils import cast_to_output_type
+from vellum.workflows.resolvers.base import BaseWorkflowResolver
+from vellum.workflows.resolvers.types import LoadStateResult
+from vellum.workflows.state.base import BaseState
+
+logger = logging.getLogger(__name__)
+
+
+class VellumResolver(BaseWorkflowResolver):
+    def get_latest_execution_events(self) -> Iterator[WorkflowEvent]:
+        return iter([])
+
+    def get_state_snapshot_history(self) -> Iterator[BaseState]:
+        return iter([])
+
+    def _deserialize_state(self, state_data: dict, state_class: Type[BaseState]) -> BaseState:
+        """Deserialize state data with proper type conversion for complex types like List[ChatMessage]."""
+        converted_data = {}
+
+        annotations = getattr(state_class, "__annotations__", {})
+
+        for field_name, field_value in state_data.items():
+            if field_name in annotations:
+                field_type = annotations[field_name]
+                converted_data[field_name] = cast_to_output_type(field_value, field_type)
+            else:
+                converted_data[field_name] = field_value
+
+        return state_class(**converted_data)
+
+    def load_state(self, previous_execution_id: Optional[Union[UUID, str]] = None) -> Optional[LoadStateResult]:
+        if isinstance(previous_execution_id, UUID):
+            previous_execution_id = str(previous_execution_id)
+
+        if previous_execution_id is None:
+            return None
+
+        if not self._context:
+            logger.warning("Cannot load state: No workflow context registered")
+            return None
+
+        client = self._context.vellum_client
+        response = client.workflows.retrieve_state(
+            span_id=previous_execution_id,
+        )
+
+        if response.state is None:
+            return None
+
+        if "meta" in response.state:
+            response.state.pop("meta")
+
+        if self._workflow_class:
+            state_class = self._workflow_class.get_state_class()
+            state = self._deserialize_state(response.state, state_class)
+        else:
+            logger.warning("No workflow class registered, falling back to BaseState")
+            state = BaseState(**response.state)
+
+        return LoadStateResult(
+            state=state,
+            previous_trace_id=response.trace_id,
+            previous_span_id=response.span_id,
+            root_trace_id=response.root_trace_id or response.trace_id,
+            root_span_id=response.root_span_id or response.span_id,
+        )

vellum/workflows/resolvers/tests/test_resolver.py

@@ -0,0 +1,168 @@
+from datetime import datetime
+from unittest.mock import Mock
+from uuid import uuid4
+from typing import List
+
+from vellum import ChatMessage
+from vellum.client.types.workflow_resolved_state import WorkflowResolvedState
+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
+
+
+def test_load_state_with_context_success():
+    """Test load_state successfully loads state when context and client are available."""
+    resolver = VellumResolver()
+    execution_id = uuid4()
+
+    class TestState(BaseState):
+        test_key: str = "test_value"
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, TestState]):
+        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": prev_id,
+            "span_id": prev_span_id,
+            "updated_ts": datetime.now().isoformat(),
+            "workflow_inputs": BaseInputs(),
+            "external_inputs": {},
+            "node_outputs": {},
+            "node_execution_cache": NodeExecutionCache(),
+            "parent": None,
+        },
+    }
+
+    previous_trace_id = str(uuid4())
+    previous_span_id = str(uuid4())
+    trace_id = str(uuid4())
+    root_trace_id = str(uuid4())
+    root_span_id = str(uuid4())
+
+    mock_response = WorkflowResolvedState(
+        trace_id=trace_id,
+        timestamp=datetime.now(),
+        span_id=str(execution_id),
+        state=state_dict,
+        previous_trace_id=previous_trace_id,
+        previous_span_id=previous_span_id,
+        root_trace_id=root_trace_id,
+        root_span_id=root_span_id,
+    )
+
+    mock_client = Mock()
+    mock_client.workflows.retrieve_state.return_value = mock_response
+
+    # 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 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 == trace_id
+    assert result.previous_span_id == str(execution_id)
+    assert result.root_trace_id == root_trace_id
+    assert result.root_span_id == root_span_id
+
+    mock_client.workflows.retrieve_state.assert_called_once_with(span_id=str(execution_id))
+
+
+def test_load_state_with_chat_message_list():
+    """Test load_state successfully loads state with chat_history containing ChatMessage list."""
+    resolver = VellumResolver()
+    execution_id = uuid4()
+
+    class TestStateWithChatHistory(BaseState):
+        test_key: str = "test_value"
+        chat_history: List[ChatMessage] = []
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, TestStateWithChatHistory]):
+        pass
+
+    # GIVEN a state dictionary with chat_history containing ChatMessage objects
+    prev_id = str(uuid4())
+    prev_span_id = str(uuid4())
+    state_dict = {
+        "test_key": "test_value",
+        "chat_history": [
+            {"role": "USER", "text": "Hello, how are you?"},
+            {"role": "ASSISTANT", "text": "I'm doing well, thank you!"},
+            {"role": "USER", "text": "What can you help me with?"},
+        ],
+        "meta": {
+            "workflow_definition": "MockWorkflow",
+            "id": prev_id,
+            "span_id": prev_span_id,
+            "updated_ts": datetime.now().isoformat(),
+            "workflow_inputs": BaseInputs(),
+            "external_inputs": {},
+            "node_outputs": {},
+            "node_execution_cache": NodeExecutionCache(),
+            "parent": None,
+        },
+    }
+
+    previous_trace_id = str(uuid4())
+    previous_span_id = str(uuid4())
+    root_trace_id = str(uuid4())
+    root_span_id = str(uuid4())
+
+    mock_response = WorkflowResolvedState(
+        trace_id=str(uuid4()),
+        timestamp=datetime.now(),
+        span_id=str(execution_id),
+        state=state_dict,
+        previous_trace_id=previous_trace_id,
+        previous_span_id=previous_span_id,
+        root_trace_id=root_trace_id,
+        root_span_id=root_span_id,
+    )
+
+    mock_client = Mock()
+    mock_client.workflows.retrieve_state.return_value = mock_response
+
+    # 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 LoadStateResult with state containing chat_history
+    assert isinstance(result, LoadStateResult)
+    assert result.state is not None
+    assert isinstance(result.state, TestStateWithChatHistory)
+    assert result.state.test_key == "test_value"
+
+    # AND the chat_history should be properly deserialized as ChatMessage objects
+    assert len(result.state.chat_history) == 3
+    assert all(isinstance(msg, ChatMessage) for msg in result.state.chat_history)
+    assert result.state.chat_history[0].role == "USER"
+    assert result.state.chat_history[0].text == "Hello, how are you?"
+    assert result.state.chat_history[1].role == "ASSISTANT"
+    assert result.state.chat_history[1].text == "I'm doing well, thank you!"
+    assert result.state.chat_history[2].role == "USER"
+    assert result.state.chat_history[2].text == "What can you help me with?"
+
+    mock_client.workflows.retrieve_state.assert_called_once_with(span_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