vellum-ai 0.9.16rc2__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

vellum/workflows/emitters/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseWorkflowEmitter
+
+__all__ = [
+    "BaseWorkflowEmitter",
+]

vellum/workflows/emitters/base.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+
+from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.state.base import BaseState
+
+
+class BaseWorkflowEmitter(ABC):
+    @abstractmethod
+    def emit_event(self, event: WorkflowEvent) -> None:
+        pass
+
+    @abstractmethod
+    def snapshot_state(self, state: BaseState) -> None:
+        pass

vellum/workflows/environment/__init__.py

@@ -0,0 +1,5 @@
+from .environment import Environment
+
+__all__ = [
+    "Environment",
+]

vellum/workflows/environment/environment.py

@@ -0,0 +1,7 @@
+from vellum.workflows.references import EnvironmentVariableReference
+
+
+class Environment:
+    @staticmethod
+    def get(name: str) -> EnvironmentVariableReference:
+        return EnvironmentVariableReference(name=name)

vellum/workflows/errors/__init__.py

@@ -0,0 +1,6 @@
+from .types import VellumError, VellumErrorCode
+
+__all__ = [
+    "VellumError",
+    "VellumErrorCode",
+]

vellum/workflows/errors/types.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from enum import Enum
+
+
+class VellumErrorCode(Enum):
+    INVALID_WORKFLOW = "INVALID_WORKFLOW"
+    INVALID_INPUTS = "INVALID_INPUTS"
+    INVALID_OUTPUTS = "INVALID_OUTPUTS"
+    INVALID_STATE = "INVALID_STATE"
+    INVALID_TEMPLATE = "INVALID_TEMPLATE"
+    INTERNAL_ERROR = "INTERNAL_ERROR"
+    PROVIDER_ERROR = "PROVIDER_ERROR"
+    USER_DEFINED_ERROR = "USER_DEFINED_ERROR"
+    WORKFLOW_CANCELLED = "WORKFLOW_CANCELLED"
+
+
+@dataclass(frozen=True)
+class VellumError:
+    message: str
+    code: VellumErrorCode

vellum/workflows/events/__init__.py

@@ -0,0 +1,31 @@
+from .node import (
+    NodeEvent,
+    NodeExecutionFulfilledEvent,
+    NodeExecutionInitiatedEvent,
+    NodeExecutionRejectedEvent,
+    NodeExecutionStreamingEvent,
+)
+from .types import WorkflowEventType
+from .workflow import (
+    WorkflowEvent,
+    WorkflowEventStream,
+    WorkflowExecutionFulfilledEvent,
+    WorkflowExecutionInitiatedEvent,
+    WorkflowExecutionRejectedEvent,
+    WorkflowExecutionStreamingEvent,
+)
+
+__all__ = [
+    "NodeExecutionFulfilledEvent",
+    "WorkflowExecutionFulfilledEvent",
+    "NodeExecutionInitiatedEvent",
+    "WorkflowExecutionInitiatedEvent",
+    "NodeEvent",
+    "NodeExecutionRejectedEvent",
+    "WorkflowExecutionRejectedEvent",
+    "NodeExecutionStreamingEvent",
+    "WorkflowExecutionStreamingEvent",
+    "WorkflowEvent",
+    "WorkflowEventStream",
+    "WorkflowEventType",
+]

vellum/workflows/events/node.py

@@ -0,0 +1,125 @@
+from typing import Any, Dict, Generic, Literal, Type, Union
+
+from pydantic import field_serializer
+
+from vellum.core.pydantic_utilities import UniversalBaseModel
+
+from vellum.workflows.errors import VellumError
+from vellum.workflows.expressions.accessor import AccessorExpression
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.references.node import NodeReference
+from vellum.workflows.types.generics import OutputsType
+
+from .types import BaseEvent, default_serializer, serialize_type_encoder
+
+
+class _BaseNodeExecutionBody(UniversalBaseModel):
+    node_definition: Type[BaseNode]
+
+    @field_serializer("node_definition")
+    def serialize_node_definition(self, node_definition: Type, _info: Any) -> Dict[str, Any]:
+        return serialize_type_encoder(node_definition)
+
+
+class _BaseNodeEvent(BaseEvent):
+    body: _BaseNodeExecutionBody
+
+    @property
+    def node_definition(self) -> Type[BaseNode]:
+        return self.body.node_definition
+
+
+NodeInputName = Union[NodeReference, AccessorExpression]
+
+
+class NodeExecutionInitiatedBody(_BaseNodeExecutionBody):
+    inputs: Dict[NodeInputName, Any]
+
+    @field_serializer("inputs")
+    def serialize_inputs(self, inputs: Dict[NodeInputName, Any], _info: Any) -> Dict[str, Any]:
+        return default_serializer({descriptor.name: value for descriptor, value in inputs.items()})
+
+
+class NodeExecutionInitiatedEvent(_BaseNodeEvent):
+    name: Literal["node.execution.initiated"] = "node.execution.initiated"
+    body: NodeExecutionInitiatedBody
+
+    @property
+    def inputs(self) -> Dict[NodeInputName, Any]:
+        return self.body.inputs
+
+
+class NodeExecutionStreamingBody(_BaseNodeExecutionBody):
+    output: BaseOutput
+
+    @field_serializer("output")
+    def serialize_output(self, output: BaseOutput, _info: Any) -> Dict[str, Any]:
+        return default_serializer(output)
+
+
+class NodeExecutionStreamingEvent(_BaseNodeEvent):
+    name: Literal["node.execution.streaming"] = "node.execution.streaming"
+    body: NodeExecutionStreamingBody
+
+    @property
+    def output(self) -> BaseOutput:
+        return self.body.output
+
+
+class NodeExecutionFulfilledBody(_BaseNodeExecutionBody, Generic[OutputsType]):
+    outputs: OutputsType
+
+    @field_serializer("outputs")
+    def serialize_outputs(self, outputs: OutputsType, _info: Any) -> Dict[str, Any]:
+        return default_serializer(outputs)
+
+
+class NodeExecutionFulfilledEvent(_BaseNodeEvent, Generic[OutputsType]):
+    name: Literal["node.execution.fulfilled"] = "node.execution.fulfilled"
+    body: NodeExecutionFulfilledBody[OutputsType]
+
+    @property
+    def outputs(self) -> OutputsType:
+        return self.body.outputs
+
+
+class NodeExecutionRejectedBody(_BaseNodeExecutionBody):
+    error: VellumError
+
+
+class NodeExecutionRejectedEvent(_BaseNodeEvent):
+    name: Literal["node.execution.rejected"] = "node.execution.rejected"
+    body: NodeExecutionRejectedBody
+
+    @property
+    def error(self) -> VellumError:
+        return self.body.error
+
+
+class NodeExecutionPausedBody(_BaseNodeExecutionBody):
+    pass
+
+
+class NodeExecutionPausedEvent(_BaseNodeEvent):
+    name: Literal["node.execution.paused"] = "node.execution.paused"
+    body: NodeExecutionPausedBody
+
+
+class NodeExecutionResumedBody(_BaseNodeExecutionBody):
+    pass
+
+
+class NodeExecutionResumedEvent(_BaseNodeEvent):
+    name: Literal["node.execution.resumed"] = "node.execution.resumed"
+    body: NodeExecutionResumedBody
+
+
+NodeEvent = Union[
+    NodeExecutionInitiatedEvent,
+    NodeExecutionStreamingEvent,
+    NodeExecutionFulfilledEvent,
+    NodeExecutionRejectedEvent,
+    NodeExecutionPausedEvent,
+    NodeExecutionResumedEvent,
+]

vellum/workflows/events/tests/test_event.py

@@ -0,0 +1,216 @@
+import pytest
+from datetime import datetime
+import json
+from uuid import UUID
+
+from deepdiff import DeepDiff
+
+from vellum.workflows.errors.types import VellumError, VellumErrorCode
+from vellum.workflows.events.node import NodeExecutionInitiatedBody, NodeExecutionInitiatedEvent
+from vellum.workflows.events.workflow import (
+    WorkflowExecutionFulfilledBody,
+    WorkflowExecutionFulfilledEvent,
+    WorkflowExecutionInitiatedBody,
+    WorkflowExecutionInitiatedEvent,
+    WorkflowExecutionRejectedBody,
+    WorkflowExecutionRejectedEvent,
+    WorkflowExecutionStreamingBody,
+    WorkflowExecutionStreamingEvent,
+)
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.workflows.base import BaseWorkflow
+
+
+class MockInputs(BaseInputs):
+    foo: str
+
+
+class MockNode(BaseNode):
+    node_foo = MockInputs.foo
+
+    class Outputs(BaseNode.Outputs):
+        example: str
+
+
+class MockWorkflow(BaseWorkflow[MockInputs, BaseState]):
+    graph = MockNode
+
+
+name_parts = __name__.split(".")
+module_root = name_parts[: name_parts.index("events")]
+
+
+@pytest.mark.parametrize(
+    ["event", "expected_json"],
+    [
+        (
+            WorkflowExecutionInitiatedEvent(
+                id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                timestamp=datetime(2024, 1, 1, 12, 0, 0),
+                trace_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                body=WorkflowExecutionInitiatedBody(
+                    workflow_definition=MockWorkflow,
+                    inputs=MockInputs(foo="bar"),
+                ),
+            ),
+            {
+                "id": "123e4567-e89b-12d3-a456-426614174000",
+                "api_version": "2024-10-25",
+                "timestamp": "2024-01-01T12:00:00",
+                "trace_id": "123e4567-e89b-12d3-a456-426614174000",
+                "span_id": "123e4567-e89b-12d3-a456-426614174000",
+                "name": "workflow.execution.initiated",
+                "body": {
+                    "workflow_definition": {
+                        "name": "MockWorkflow",
+                        "module": module_root + ["events", "tests", "test_event"],
+                    },
+                    "inputs": {
+                        "foo": "bar",
+                    },
+                },
+            },
+        ),
+        (
+            NodeExecutionInitiatedEvent(
+                id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                timestamp=datetime(2024, 1, 1, 12, 0, 0),
+                trace_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                body=NodeExecutionInitiatedBody(
+                    node_definition=MockNode,
+                    inputs={
+                        MockNode.node_foo: "bar",
+                    },
+                ),
+            ),
+            {
+                "id": "123e4567-e89b-12d3-a456-426614174000",
+                "api_version": "2024-10-25",
+                "timestamp": "2024-01-01T12:00:00",
+                "trace_id": "123e4567-e89b-12d3-a456-426614174000",
+                "span_id": "123e4567-e89b-12d3-a456-426614174000",
+                "name": "node.execution.initiated",
+                "body": {
+                    "node_definition": {
+                        "name": "MockNode",
+                        "module": module_root + ["events", "tests", "test_event"],
+                    },
+                    "inputs": {
+                        "node_foo": "bar",
+                    },
+                },
+            },
+        ),
+        (
+            WorkflowExecutionStreamingEvent(
+                id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                timestamp=datetime(2024, 1, 1, 12, 0, 0),
+                trace_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                body=WorkflowExecutionStreamingBody(
+                    workflow_definition=MockWorkflow,
+                    output=BaseOutput(
+                        name="example",
+                        value="foo",
+                    ),
+                ),
+            ),
+            {
+                "id": "123e4567-e89b-12d3-a456-426614174000",
+                "api_version": "2024-10-25",
+                "timestamp": "2024-01-01T12:00:00",
+                "trace_id": "123e4567-e89b-12d3-a456-426614174000",
+                "span_id": "123e4567-e89b-12d3-a456-426614174000",
+                "name": "workflow.execution.streaming",
+                "body": {
+                    "workflow_definition": {
+                        "name": "MockWorkflow",
+                        "module": module_root + ["events", "tests", "test_event"],
+                    },
+                    "output": {
+                        "name": "example",
+                        "value": "foo",
+                    },
+                },
+            },
+        ),
+        (
+            WorkflowExecutionFulfilledEvent(
+                id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                timestamp=datetime(2024, 1, 1, 12, 0, 0),
+                trace_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                body=WorkflowExecutionFulfilledBody(
+                    workflow_definition=MockWorkflow,
+                    outputs=MockNode.Outputs(
+                        example="foo",
+                    ),
+                ),
+            ),
+            {
+                "id": "123e4567-e89b-12d3-a456-426614174000",
+                "api_version": "2024-10-25",
+                "timestamp": "2024-01-01T12:00:00",
+                "trace_id": "123e4567-e89b-12d3-a456-426614174000",
+                "span_id": "123e4567-e89b-12d3-a456-426614174000",
+                "name": "workflow.execution.fulfilled",
+                "body": {
+                    "workflow_definition": {
+                        "name": "MockWorkflow",
+                        "module": module_root + ["events", "tests", "test_event"],
+                    },
+                    "outputs": {
+                        "example": "foo",
+                    },
+                },
+            },
+        ),
+        (
+            WorkflowExecutionRejectedEvent(
+                id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                timestamp=datetime(2024, 1, 1, 12, 0, 0),
+                trace_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+                body=WorkflowExecutionRejectedBody(
+                    workflow_definition=MockWorkflow,
+                    error=VellumError(
+                        message="Workflow failed",
+                        code=VellumErrorCode.USER_DEFINED_ERROR,
+                    ),
+                ),
+            ),
+            {
+                "id": "123e4567-e89b-12d3-a456-426614174000",
+                "api_version": "2024-10-25",
+                "timestamp": "2024-01-01T12:00:00",
+                "trace_id": "123e4567-e89b-12d3-a456-426614174000",
+                "span_id": "123e4567-e89b-12d3-a456-426614174000",
+                "name": "workflow.execution.rejected",
+                "body": {
+                    "workflow_definition": {
+                        "name": "MockWorkflow",
+                        "module": module_root + ["events", "tests", "test_event"],
+                    },
+                    "error": {
+                        "message": "Workflow failed",
+                        "code": "USER_DEFINED_ERROR",
+                    },
+                },
+            },
+        ),
+    ],
+    ids=[
+        "workflow.execution.initiated",
+        "node.execution.initiated",
+        "workflow.execution.streaming",
+        "workflow.execution.fulfilled",
+        "workflow.execution.rejected",
+    ],
+)
+def test_event_serialization(event, expected_json):
+    assert not DeepDiff(json.loads(event.model_dump_json()), expected_json)

vellum/workflows/events/types.py

@@ -0,0 +1,52 @@
+from datetime import datetime
+from enum import Enum
+import json
+from uuid import UUID, uuid4
+from typing import Any, Dict, Literal
+
+from pydantic import Field
+
+from vellum.core.pydantic_utilities import UniversalBaseModel
+
+from vellum.workflows.state.encoder import DefaultStateEncoder
+from vellum.workflows.types.utils import datetime_now
+
+
+class WorkflowEventType(Enum):
+    NODE = "NODE"
+    WORKFLOW = "WORKFLOW"
+
+
+def default_datetime_factory() -> datetime:
+    """
+    Makes it possible to mock the datetime factory for testing.
+    """
+
+    return datetime_now()
+
+
+excluded_modules = {"typing", "builtins"}
+
+
+def serialize_type_encoder(obj: type) -> Dict[str, Any]:
+    return {
+        "name": obj.__name__,
+        "module": obj.__module__.split("."),
+    }
+
+
+def default_serializer(obj: Any) -> Any:
+    return json.loads(
+        json.dumps(
+            obj,
+            cls=DefaultStateEncoder,
+        )
+    )
+
+
+class BaseEvent(UniversalBaseModel):
+    id: UUID = Field(default_factory=uuid4)
+    timestamp: datetime = Field(default_factory=default_datetime_factory)
+    api_version: Literal["2024-10-25"] = "2024-10-25"
+    trace_id: UUID
+    span_id: UUID

vellum/workflows/events/utils.py

@@ -0,0 +1,5 @@
+from vellum.workflows.events.workflow import WorkflowEvent
+
+
+def is_terminal_event(event: WorkflowEvent) -> bool:
+    return event.name in {"workflow.execution.fulfilled", "workflow.execution.rejected", "workflow.execution.paused"}

vellum/workflows/events/workflow.py

@@ -0,0 +1,139 @@
+from typing import TYPE_CHECKING, Any, Dict, Generator, Generic, Iterable, Literal, Type, Union
+
+from pydantic import field_serializer
+
+from vellum.core.pydantic_utilities import UniversalBaseModel
+
+from vellum.workflows.errors import VellumError
+from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.references import ExternalInputReference
+from vellum.workflows.types.generics import OutputsType, WorkflowInputsType
+
+from .node import (
+    NodeExecutionFulfilledEvent,
+    NodeExecutionInitiatedEvent,
+    NodeExecutionPausedEvent,
+    NodeExecutionRejectedEvent,
+    NodeExecutionResumedEvent,
+    NodeExecutionStreamingEvent,
+)
+from .types import BaseEvent, default_serializer, serialize_type_encoder
+
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+
+class _BaseWorkflowExecutionBody(UniversalBaseModel):
+    workflow_definition: Type["BaseWorkflow"]
+
+    @field_serializer("workflow_definition")
+    def serialize_workflow_definition(self, workflow_definition: Type, _info: Any) -> Dict[str, Any]:
+        return serialize_type_encoder(workflow_definition)
+
+
+class WorkflowExecutionInitiatedBody(_BaseWorkflowExecutionBody, Generic[WorkflowInputsType]):
+    inputs: WorkflowInputsType
+
+    @field_serializer("inputs")
+    def serialize_inputs(self, inputs: WorkflowInputsType, _info: Any) -> Dict[str, Any]:
+        return default_serializer(inputs)
+
+
+class WorkflowExecutionInitiatedEvent(BaseEvent, Generic[WorkflowInputsType]):
+    name: Literal["workflow.execution.initiated"] = "workflow.execution.initiated"
+    body: WorkflowExecutionInitiatedBody[WorkflowInputsType]
+
+    @property
+    def inputs(self) -> WorkflowInputsType:
+        return self.body.inputs
+
+
+class WorkflowExecutionStreamingBody(_BaseWorkflowExecutionBody):
+    output: BaseOutput
+
+    @field_serializer("output")
+    def serialize_output(self, output: BaseOutput, _info: Any) -> Dict[str, Any]:
+        return default_serializer(output)
+
+
+class WorkflowExecutionStreamingEvent(BaseEvent):
+    name: Literal["workflow.execution.streaming"] = "workflow.execution.streaming"
+    body: WorkflowExecutionStreamingBody
+
+    @property
+    def output(self) -> BaseOutput:
+        return self.body.output
+
+
+class WorkflowExecutionFulfilledBody(_BaseWorkflowExecutionBody, Generic[OutputsType]):
+    outputs: OutputsType
+
+    @field_serializer("outputs")
+    def serialize_outputs(self, outputs: OutputsType, _info: Any) -> Dict[str, Any]:
+        return default_serializer(outputs)
+
+
+class WorkflowExecutionFulfilledEvent(BaseEvent, Generic[OutputsType]):
+    name: Literal["workflow.execution.fulfilled"] = "workflow.execution.fulfilled"
+    body: WorkflowExecutionFulfilledBody[OutputsType]
+
+    @property
+    def outputs(self) -> OutputsType:
+        return self.body.outputs
+
+
+class WorkflowExecutionRejectedBody(_BaseWorkflowExecutionBody):
+    error: VellumError
+
+
+class WorkflowExecutionRejectedEvent(BaseEvent):
+    name: Literal["workflow.execution.rejected"] = "workflow.execution.rejected"
+    body: WorkflowExecutionRejectedBody
+
+    @property
+    def error(self) -> VellumError:
+        return self.body.error
+
+
+class WorkflowExecutionPausedBody(_BaseWorkflowExecutionBody):
+    external_inputs: Iterable[ExternalInputReference]
+
+
+class WorkflowExecutionPausedEvent(BaseEvent):
+    name: Literal["workflow.execution.paused"] = "workflow.execution.paused"
+    body: WorkflowExecutionPausedBody
+
+    @property
+    def external_inputs(self) -> Iterable[ExternalInputReference]:
+        return self.body.external_inputs
+
+
+class WorkflowExecutionResumedBody(_BaseWorkflowExecutionBody):
+    pass
+
+
+class WorkflowExecutionResumedEvent(BaseEvent):
+    name: Literal["workflow.execution.resumed"] = "workflow.execution.resumed"
+    body: WorkflowExecutionResumedBody
+
+
+GenericWorkflowEvent = Union[
+    WorkflowExecutionStreamingEvent,
+    WorkflowExecutionRejectedEvent,
+    WorkflowExecutionPausedEvent,
+    WorkflowExecutionResumedEvent,
+    NodeExecutionInitiatedEvent,
+    NodeExecutionStreamingEvent,
+    NodeExecutionFulfilledEvent,
+    NodeExecutionRejectedEvent,
+    NodeExecutionPausedEvent,
+    NodeExecutionResumedEvent,
+]
+
+WorkflowEvent = Union[
+    GenericWorkflowEvent,
+    WorkflowExecutionInitiatedEvent,
+    WorkflowExecutionFulfilledEvent,
+]
+
+WorkflowEventStream = Generator[WorkflowEvent, None, None]

vellum/workflows/exceptions.py

@@ -0,0 +1,15 @@
+from vellum.workflows.errors import VellumError, VellumErrorCode
+
+
+class NodeException(Exception):
+    def __init__(self, message: str, code: VellumErrorCode):
+        self.message = message
+        self.code = code
+        super().__init__(message)
+
+    @property
+    def error(self) -> VellumError:
+        return VellumError(
+            message=self.message,
+            code=self.code,
+        )