vellum-ai 0.9.16rc2__py3-none-any.whl → 0.10.0__py3-none-any.whl

vellum/workflows/references/output.py

@@ -0,0 +1,78 @@
+from queue import Queue
+from typing import TYPE_CHECKING, Any, Generator, Generic, Optional, Tuple, Type, TypeVar, cast
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+from vellum.workflows.constants import UNDEF
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+if TYPE_CHECKING:
+    from vellum.workflows.outputs import BaseOutputs
+    from vellum.workflows.state.base import BaseState
+
+
+_OutputType = TypeVar("_OutputType")
+
+
+class OutputReference(BaseDescriptor[_OutputType], Generic[_OutputType]):
+    def __init__(
+        self,
+        *,
+        name: str,
+        types: Tuple[Type[_OutputType], ...],
+        instance: Optional[_OutputType],
+        outputs_class: Type["BaseOutputs"],
+    ) -> None:
+        super().__init__(name=name, types=types, instance=instance)
+        self._outputs_class = outputs_class
+
+    @property
+    def outputs_class(self) -> Type["BaseOutputs"]:
+        return self._outputs_class
+
+    def resolve(self, state: "BaseState") -> _OutputType:
+        node_output = state.meta.node_outputs.get(self, UNDEF)
+        if isinstance(node_output, Queue):
+            # Fix typing surrounding the return value of node outputs
+            # https://app.shortcut.com/vellum/story/4783
+            return self._as_generator(node_output)  # type: ignore[return-value]
+
+        if node_output is not UNDEF:
+            return cast(_OutputType, node_output)
+
+        if state.meta.parent:
+            return self.resolve(state.meta.parent)
+
+        # Fix typing surrounding the return value of node outputs
+        # https://app.shortcut.com/vellum/story/4783
+        return cast(Type[UNDEF], node_output)  # type: ignore[return-value]
+
+    def _as_generator(self, node_output: Queue) -> Generator[_OutputType, None, Type[UNDEF]]:
+        while True:
+            item = node_output.get()
+            if item is UNDEF:
+                return UNDEF
+            yield cast(_OutputType, item)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, type(self)):
+            return False
+        return super().__eq__(other) and self._outputs_class == other._outputs_class
+
+    def __hash__(self) -> int:
+        return hash((self._outputs_class, self._name))
+
+    def __repr__(self) -> str:
+        return f"{self._outputs_class.__qualname__}.{self.name}"
+
+    def __deepcopy__(self, memo: dict) -> "OutputReference[_OutputType]":
+        return OutputReference(
+            name=self._name, types=self._types, instance=self._instance, outputs_class=self._outputs_class
+        )
+
+    @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/state_value.py

@@ -0,0 +1,23 @@
+from typing import TYPE_CHECKING, TypeVar, cast
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+
+
+_T = TypeVar("_T")
+
+
+class StateValueReference(BaseDescriptor[_T]):
+
+    def resolve(self, state: "BaseState") -> _T:
+        if hasattr(state, self._name):
+            return cast(_T, getattr(state, self._name))
+
+        if state.meta.parent:
+            return self.resolve(state.meta.parent)
+
+        raise NodeException(f"Missing required Workflow state: {self._name}", code=VellumErrorCode.INVALID_STATE)

vellum/workflows/references/vellum_secret.py

@@ -0,0 +1,15 @@
+from typing import TYPE_CHECKING
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.types.core import VellumSecret
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+
+
+class VellumSecretReference(BaseDescriptor[VellumSecret]):
+    def __init__(self, name: str):
+        super().__init__(name=name, types=(VellumSecret,))
+
+    def resolve(self, state: "BaseState") -> VellumSecret:
+        return VellumSecret(name=self.name)

vellum/workflows/references/workflow_input.py

@@ -0,0 +1,41 @@
+from typing import TYPE_CHECKING, Generic, Optional, Tuple, Type, TypeVar, cast
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+
+if TYPE_CHECKING:
+    from vellum.workflows.inputs.base import BaseInputs
+    from vellum.workflows.state.base import BaseState
+
+_InputType = TypeVar("_InputType")
+
+
+class WorkflowInputReference(BaseDescriptor[_InputType], Generic[_InputType]):
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        types: Tuple[Type[_InputType], ...],
+        instance: Optional[_InputType],
+        inputs_class: Type["BaseInputs"],
+    ) -> None:
+        super().__init__(name=name, types=types, instance=instance)
+        self._inputs_class = inputs_class
+
+    @property
+    def inputs_class(self) -> Type["BaseInputs"]:
+        return self._inputs_class
+
+    def resolve(self, state: "BaseState") -> _InputType:
+        if hasattr(state.meta.workflow_inputs, self._name):
+            return cast(_InputType, getattr(state.meta.workflow_inputs, self._name))
+
+        if state.meta.parent:
+            return self.resolve(state.meta.parent)
+
+        raise NodeException(f"Missing required Workflow input: {self._name}", code=VellumErrorCode.INVALID_INPUTS)
+
+    def __repr__(self) -> str:
+        return f"{self._inputs_class.__qualname__}.{self.name}"

vellum/workflows/resolvers/__init__.py

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

vellum/workflows/resolvers/base.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+from typing import Iterator
+
+from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.state.base import BaseState
+
+
+class BaseWorkflowResolver(ABC):
+    @abstractmethod
+    def get_latest_execution_events(self) -> Iterator[WorkflowEvent]:
+        pass
+
+    @abstractmethod
+    def get_state_snapshot_history(self) -> Iterator[BaseState]:
+        pass

vellum/workflows/runner/__init__.py

@@ -0,0 +1,5 @@
+from .runner import WorkflowRunner
+
+__all__ = [
+    "WorkflowRunner",
+]