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

vellum/workflows/outputs/base.py

@@ -0,0 +1,196 @@
+from typing import TYPE_CHECKING, Any, Generic, Iterator, Set, Tuple, Type, TypeVar, Union, cast
+from typing_extensions import dataclass_transform
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+from vellum.workflows.constants import UNDEF
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.types.utils import get_class_attr_names, infer_types
+
+if TYPE_CHECKING:
+    from vellum.workflows.nodes.bases.base import BaseNode
+
+_Delta = TypeVar("_Delta")
+_Accumulated = TypeVar("_Accumulated")
+
+
+class BaseOutput(Generic[_Delta, _Accumulated]):
+    _value: Union[_Accumulated, Type[UNDEF]]
+    _delta: Union[_Delta, Type[UNDEF]]
+    _name: str
+
+    def __init__(
+        self,
+        name: str,
+        value: Union[_Accumulated, Type[UNDEF]] = UNDEF,
+        delta: Union[_Delta, Type[UNDEF]] = UNDEF,
+    ) -> None:
+        if value is not UNDEF and delta is not UNDEF:
+            raise ValueError("Cannot set both value and delta")
+
+        self._name = name
+        self._value = value
+        self._delta = delta
+
+    @property
+    def delta(self) -> Union[_Delta, Type[UNDEF]]:
+        return self._delta
+
+    @property
+    def value(self) -> Union[_Accumulated, Type[UNDEF]]:
+        return self._value
+
+    @property
+    def is_initiated(self) -> bool:
+        return self._delta is UNDEF and self._value is UNDEF
+
+    @property
+    def is_streaming(self) -> bool:
+        return self._delta is not UNDEF and self._value is UNDEF
+
+    @property
+    def is_fulfilled(self) -> bool:
+        return self._delta is UNDEF and self._value is not UNDEF
+
+    @property
+    def name(self) -> str:
+        return 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)
+
+    def serialize(self) -> dict:
+        data: dict[str, Any] = {
+            "name": self.name,
+        }
+
+        if self.value is not UNDEF:
+            data["value"] = self.value
+
+        if self.delta is not UNDEF:
+            data["delta"] = self.delta
+
+        return data
+
+
+@dataclass_transform(kw_only_default=True)
+class _BaseOutputsMeta(type):
+    def __eq__(cls, other: Any) -> bool:
+        """
+        We need to include custom eq logic to prevent infinite loops during ipython reloading.
+        """
+
+        if not isinstance(other, _BaseOutputsMeta):
+            return False
+
+        if not cls.__qualname__.endswith(".Outputs") or not other.__qualname__.endswith(".Outputs"):
+            return super().__eq__(other)
+
+        self_outputs_class = cast(Type["BaseNode.Outputs"], cls)
+        other_outputs_class = cast(Type["BaseNode.Outputs"], other)
+
+        if not hasattr(self_outputs_class, "_node_class") or not hasattr(other_outputs_class, "_node_class"):
+            return super().__eq__(other)
+
+        if self_outputs_class._node_class is None or other_outputs_class._node_class is None:
+            return super().__eq__(other)
+
+        return getattr(self_outputs_class._node_class, "__qualname__") == getattr(
+            other_outputs_class._node_class, "__qualname__"
+        )
+
+    def __setattr__(cls, name: str, value: Any) -> None:
+        if isinstance(value, OutputReference):
+            # During module reload, tools like ipython will set class attributes from the get attribute
+            # from the old class. Our getattribute dynamically returns OutputDescriptors, which means during
+            # module reload, it will override the class' original default value with its own OutputReference.
+            # We want to avoid this, so we check if the name of the class and the name of the descriptor
+            # are the same, and if they are, we don't set the attribute.
+            if f"{cls.__qualname__}.{name}" == str(value):
+                return super().__setattr__(name, value.instance)
+
+        return super().__setattr__(name, value)
+
+    def __getattribute__(cls, name: str) -> Any:
+        if name.startswith("_") or not issubclass(cls, BaseOutputs):
+            return super().__getattribute__(name)
+
+        attr_names = get_class_attr_names(cls)
+        if name in attr_names:
+            # We first try to resolve the instance that this class attribute name is mapped to. If it's not found,
+            # we iterate through its inheritance hierarchy to find the first base class that has this attribute
+            # and use its mapping.
+            instance = vars(cls).get(name, UNDEF)
+            if not instance:
+                for base in cls.__mro__[1:]:
+                    if hasattr(base, name):
+                        instance = getattr(base, name)
+                        break
+
+            types = infer_types(cls, name)
+            return OutputReference(
+                name=name,
+                types=types,
+                instance=instance,
+                outputs_class=cls,
+            )
+
+        return super().__getattribute__(name)
+
+    def __hash__(self) -> int:
+        return hash(self.__qualname__)
+
+    def __iter__(cls) -> Iterator[OutputReference]:
+        # We iterate through the inheritance hierarchy to find all the OutputDescriptors attached to this Outputs class.
+        # __mro__ is the method resolution order, which is the order in which base classes are resolved.
+        yielded_attr_names: Set[str] = set()
+
+        for resolved_cls in cls.__mro__:
+            attr_names = get_class_attr_names(resolved_cls)
+            for attr_name in attr_names:
+                if attr_name in yielded_attr_names:
+                    continue
+
+                attr_value = getattr(resolved_cls, attr_name)
+                if not isinstance(attr_value, OutputReference):
+                    continue
+
+                yield attr_value
+                yielded_attr_names.add(attr_name)
+
+
+class BaseOutputs(metaclass=_BaseOutputsMeta):
+    def __init__(self, **kwargs: Any) -> None:
+        for name, value in kwargs.items():
+            setattr(self, name, value)
+
+        if hasattr(self, "_outputs_post_init") and callable(self._outputs_post_init):
+            self._outputs_post_init(**kwargs)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, dict):
+            return super().__eq__(other)
+
+        outputs = {name: value for name, value in vars(self).items() if not name.startswith("_") and value is not UNDEF}
+        return outputs == other
+
+    def __repr__(self) -> str:
+        values = f"{', '.join(f'{k}={v}' for k, v in vars(self).items() if not k.startswith('_'))}"
+        return f"{self.__class__.__name__}({values})"
+
+    def __iter__(self) -> Iterator[Tuple[OutputReference, Any]]:
+        for output_descriptor in self.__class__:
+            yield (output_descriptor, getattr(self, output_descriptor.name, output_descriptor.instance))
+
+    def __getitem__(self, key: str) -> Any:
+        return getattr(self, key)
+
+    @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/ports/__init__.py

@@ -0,0 +1,7 @@
+from .node_ports import NodePorts
+from .port import Port
+
+__all__ = [
+    "NodePorts",
+    "Port",
+]

vellum/workflows/ports/node_ports.py

@@ -0,0 +1,75 @@
+from typing import Any, Dict, Iterable, Iterator, Optional, Set, Tuple, Type
+
+from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.ports.port import Port
+from vellum.workflows.ports.utils import validate_ports
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import ConditionType
+
+
+class _NodePortsMeta(type):
+    def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        for k, v in dct.items():
+            if k.startswith("_"):
+                continue
+            if not isinstance(v, Port):
+                raise ValueError(f"All fields in {name} must be of type Port. Received: {v.__class__}")
+
+        return super().__new__(mcs, name, bases, dct)
+
+    def __iter__(cls) -> Iterator[Port]:
+        for attr_name, attr_value in cls.__dict__.items():
+            if not attr_name.startswith("_") and isinstance(attr_value, Port):
+                yield attr_value
+
+    @property
+    def _default_port(cls) -> Optional[Port]:
+        default_ports = [port for port in cls if port.default]
+
+        if len(default_ports) > 1:
+            raise ValueError(f"Class {cls.__name__} must have only one default port")
+
+        return default_ports[0] if default_ports else None
+
+
+class NodePorts(metaclass=_NodePortsMeta):
+    def __call__(self, outputs: BaseOutputs, state: BaseState) -> Iterable[Port]:
+        """
+        Invokes the appropriate ports based on the fulfilled outputs and state.
+        """
+
+        invoked_ports: Set[Port] = set()
+        all_ports = [port for port in self.__class__]
+        enforce_single_invoked_conditional_port = validate_ports(all_ports)
+
+        for port in all_ports:
+            if port._condition_type == ConditionType.IF:
+                resolved_condition = port.resolve_condition(state)
+                if resolved_condition:
+                    invoked_ports.add(port)
+                    if enforce_single_invoked_conditional_port:
+                        break
+
+            elif port._condition_type == ConditionType.ELIF:
+                resolved_condition = port.resolve_condition(state)
+                if resolved_condition:
+                    invoked_ports.add(port)
+                    break
+
+            elif port._condition_type == ConditionType.ELSE and not invoked_ports:
+                invoked_ports.add(port)
+                break
+
+        if not invoked_ports:
+            default_port = self.__class__._default_port
+            if default_port:
+                invoked_ports.add(default_port)
+
+        return invoked_ports
+
+    def __lt__(self, output: BaseOutput) -> Iterable[Port]:
+        """
+        Invokes the appropriate ports based on the streamed output
+        """
+
+        return set()

vellum/workflows/ports/port.py

@@ -0,0 +1,75 @@
+from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Type
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.edges.edge import Edge
+from vellum.workflows.graph import Graph, GraphTarget
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import ConditionType
+
+if TYPE_CHECKING:
+    from vellum.workflows.nodes.bases import BaseNode
+
+
+class Port:
+    node_class: Type["BaseNode"]
+
+    _edges: List[Edge]
+    _condition: Optional[BaseDescriptor]
+    _condition_type: Optional[ConditionType]
+
+    def __init__(
+        self,
+        default: bool = False,
+        fork_state: bool = False,
+        condition: Optional[Any] = None,
+        condition_type: Optional[ConditionType] = None,
+    ):
+        self.default = default
+        self.node_class = None  # type: ignore[assignment]
+        self._fork_state = fork_state
+        self._edges = []
+        self._condition: Optional[BaseDescriptor] = condition
+        self._condition_type: Optional[ConditionType] = condition_type
+
+    def __set_name__(self, owner: Type, name: str) -> None:
+        self.name = name
+
+    def __repr__(self) -> str:
+        return f"{self.node_class}.Ports.{self.name}"
+
+    @property
+    def fork_state(self) -> bool:
+        return self._fork_state
+
+    @property
+    def edges(self) -> Iterator[Edge]:
+        return iter(self._edges)
+
+    def __rshift__(self, other: GraphTarget) -> Graph:
+        if isinstance(other, set) or isinstance(other, Graph):
+            return Graph.from_port(self) >> other
+
+        edge = Edge(from_port=self, to_node=other)
+        if edge not in self._edges:
+            self._edges.append(edge)
+
+        return Graph.from_edge(edge)
+
+    @staticmethod
+    def on_if(condition: BaseDescriptor, fork_state: bool = False) -> "Port":
+        return Port(condition=condition, condition_type=ConditionType.IF, fork_state=fork_state)
+
+    @staticmethod
+    def on_elif(condition: BaseDescriptor, fork_state: bool = False) -> "Port":
+        return Port(condition=condition, condition_type=ConditionType.ELIF, fork_state=fork_state)
+
+    @staticmethod
+    def on_else(fork_state: bool = False) -> "Port":
+        return Port(condition_type=ConditionType.ELSE, fork_state=fork_state)
+
+    def resolve_condition(self, state: BaseState) -> bool:
+        if self._condition is None:
+            return False
+
+        value = self._condition.resolve(state)
+        return bool(value)

vellum/workflows/ports/utils.py

@@ -0,0 +1,40 @@
+from collections import Counter
+from typing import List
+
+from vellum.workflows.ports.port import Port
+from vellum.workflows.types.core import ConditionType
+
+PORT_TYPE_PRIORITIES = {
+    ConditionType.IF: 1,
+    ConditionType.ELIF: 2,
+    ConditionType.ELSE: 3,
+}
+
+
+def validate_ports(ports: List[Port]) -> bool:
+    # We don't want to validate ports with no condition (default ports)
+    port_types = [port._condition_type for port in ports if port._condition_type is not None]
+    sorted_port_types = sorted(port_types, key=lambda port_type: PORT_TYPE_PRIORITIES[port_type])
+
+    if sorted_port_types != port_types:
+        raise ValueError("Port conditions must be in the following order: on_if, on_elif, on_else")
+
+    counter = Counter(port_types)
+    number_of_if_ports = counter[ConditionType.IF]
+    number_of_elif_ports = counter[ConditionType.ELIF]
+    number_of_else_ports = counter[ConditionType.ELSE]
+    ports_class = f"{ports[0].node_class}.Ports"
+
+    if number_of_if_ports == 0 and (number_of_elif_ports > 0 or number_of_else_ports > 0):
+        raise ValueError(
+            f"Class {ports_class} containing on_elif or on_else port conditions must have at least one on_if condition"
+        )
+
+    if number_of_elif_ports > 0 and number_of_if_ports != 1:
+        raise ValueError(f"Class {ports_class} containing on_elif ports must have exactly one on_if condition")
+
+    if number_of_else_ports > 1:
+        raise ValueError(f"Class {ports_class} must have at most one on_else port condition")
+
+    enforce_single_invoked_conditional_port = number_of_elif_ports > 0 or number_of_if_ports <= 1
+    return enforce_single_invoked_conditional_port

vellum/workflows/references/__init__.py

@@ -0,0 +1,17 @@
+from .environment_variable import EnvironmentVariableReference
+from .external_input import ExternalInputReference
+from .lazy import LazyReference
+from .node import NodeReference
+from .output import OutputReference
+from .state_value import StateValueReference
+from .workflow_input import WorkflowInputReference
+
+__all__ = [
+    "EnvironmentVariableReference",
+    "ExternalInputReference",
+    "LazyReference",
+    "NodeReference",
+    "OutputReference",
+    "StateValueReference",
+    "WorkflowInputReference",
+]

vellum/workflows/references/environment_variable.py

@@ -0,0 +1,20 @@
+import os
+from typing import TYPE_CHECKING
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+
+
+class EnvironmentVariableReference(BaseDescriptor[str]):
+    def __init__(self, *, name: str) -> None:
+        super().__init__(name=name, types=(str,))
+
+    def resolve(self, state: "BaseState") -> str:
+        env_value = os.environ.get(self.name)
+        if env_value is not None:
+            return env_value
+
+        # Fetch Vellum Environment Variable named `self.name` once that project is done
+        raise ValueError(f"No environment variable named '{self.name}' found")

vellum/workflows/references/execution_count.py

@@ -0,0 +1,20 @@
+from typing import TYPE_CHECKING, Type
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+if TYPE_CHECKING:
+    from vellum.workflows.nodes.bases import BaseNode
+    from vellum.workflows.state.base import BaseState
+
+
+class ExecutionCountReference(BaseDescriptor[int]):
+
+    def __init__(
+        self,
+        node_class: Type["BaseNode"],
+    ) -> None:
+        super().__init__(name=f"Execution Count({node_class.__name__})", types=(int,))
+        self._node_class = node_class
+
+    def resolve(self, state: "BaseState") -> int:
+        return state.meta.node_execution_cache.get_execution_count(self._node_class)

vellum/workflows/references/external_input.py

@@ -0,0 +1,49 @@
+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.constants import UNDEF
+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 ExternalInputReference(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:
+        external_input = state.meta.external_inputs.get(self)
+        if external_input is not UNDEF:
+            return cast(_InputType, external_input)
+
+        if state.meta.parent:
+            return self.resolve(state.meta.parent)
+
+        raise NodeException(f"Missing required Node Input: {self._name}", code=VellumErrorCode.INVALID_INPUTS)
+
+    @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/input.py

@@ -0,0 +1,7 @@
+from typing import Union
+
+from vellum.workflows.references.external_input import ExternalInputReference
+from vellum.workflows.references.workflow_input import WorkflowInputReference
+
+# Note: We may want to consolidate these two into a single descriptor.
+InputReference = Union[ExternalInputReference, WorkflowInputReference]

vellum/workflows/references/lazy.py

@@ -0,0 +1,55 @@
+import ast
+import inspect
+from typing import TYPE_CHECKING, Callable, Generic, TypeVar, get_args
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+
+_T = TypeVar("_T")
+
+
+class LazyReference(BaseDescriptor[_T], Generic[_T]):
+    def __init__(
+        self,
+        get: Callable[[], BaseDescriptor[_T]],
+    ) -> None:
+        self._get = get
+        # TODO: figure out this some times returns empty
+        # Original example: https://github.com/vellum-ai/workflows-as-code-runner-prototype/pull/128/files#diff-67aaa468aa37b6130756bfaf93f03954d7b518617922efb3350882ea4ae03d60R36 # noqa: E501
+        # https://app.shortcut.com/vellum/story/4993
+        types = get_args(type(self))
+        super().__init__(name=self._get_name(), types=types)
+
+    def resolve(self, state: "BaseState") -> _T:
+        from vellum.workflows.descriptors.utils import resolve_value
+
+        return resolve_value(self._get(), state)
+
+    def _get_name(self) -> str:
+        """
+        We try our best to parse out the source code that generates the descriptor,
+        setting that as the descriptor's name. Names are only used for debugging, so
+        we could flesh out edge cases over time.
+        """
+        source = inspect.getsource(self._get).strip()
+        try:
+            parsed = ast.parse(source)
+            assignment = parsed.body[0]
+
+            if not isinstance(assignment, ast.Assign):
+                return source
+
+            call = assignment.value
+            if not isinstance(call, ast.Call) or not call.args:
+                return source
+
+            lambda_expression = call.args[0]
+            if not isinstance(lambda_expression, ast.Lambda):
+                return source
+
+            body = lambda_expression.body
+            return source[body.col_offset : body.end_col_offset]
+        except Exception:
+            return source

vellum/workflows/references/node.py

@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Type, TypeVar
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+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.nodes.bases.base import BaseNode
+    from vellum.workflows.state.base import BaseState
+
+_T = TypeVar("_T")
+
+
+class NodeReference(BaseDescriptor[_T]):
+    def __init__(
+        self, *, name: str, types: Tuple[Type[_T], ...], instance: Optional[_T] = None, node_class: Type["BaseNode"]
+    ) -> None:
+        self._name = name
+        self._types = types
+        self._instance = instance
+        self._node_class = node_class
+
+    @property
+    def node_class(self) -> Type["BaseNode"]:
+        return self._node_class
+
+    def resolve(self, state: "BaseState") -> _T:
+        raise NodeException(
+            f"NodeDescriptors cannot be resolved during runtime. Got: {self._name}",
+            code=VellumErrorCode.INTERNAL_ERROR,
+        )
+
+    def __repr__(self) -> str:
+        return f"{self._node_class.__qualname__}.{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)