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

vellum/workflows/expressions/not_in.py

@@ -0,0 +1,31 @@
+from typing import Generic, TypeVar, Union
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.utils import resolve_value
+from vellum.workflows.state.base import BaseState
+
+LHS = TypeVar("LHS")
+RHS = TypeVar("RHS")
+
+
+class NotInExpression(BaseDescriptor[bool], Generic[LHS, RHS]):
+    def __init__(
+        self,
+        *,
+        lhs: Union[BaseDescriptor[LHS], LHS],
+        rhs: Union[BaseDescriptor[RHS], RHS],
+    ) -> None:
+        super().__init__(name=f"{lhs} not in {rhs}", types=(bool,))
+        self._lhs = lhs
+        self._rhs = rhs
+
+    def resolve(self, state: "BaseState") -> bool:
+        # Support any type that implements the not in operator
+        # https://app.shortcut.com/vellum/story/4658
+        lhs = resolve_value(self._lhs, state)
+
+        rhs = resolve_value(self._rhs, state)
+        if not isinstance(rhs, (list, tuple, set, dict, str)):
+            raise ValueError(f"Expected a RHS that supported contains, got: {rhs.__class__.__name__}")
+
+        return lhs not in rhs

vellum/workflows/expressions/or_.py

@@ -0,0 +1,32 @@
+from typing import TypeVar, Union
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.utils import resolve_value
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.utils import resolve_combined_types
+
+LHS = TypeVar("LHS")
+RHS = TypeVar("RHS")
+
+
+class OrExpression(BaseDescriptor[Union[LHS, RHS]]):
+    def __init__(
+        self,
+        *,
+        lhs: Union[BaseDescriptor[LHS], LHS],
+        rhs: Union[BaseDescriptor[RHS], RHS],
+    ) -> None:
+        super().__init__(
+            name=f"{lhs} or {rhs}",
+            types=resolve_combined_types(lhs, rhs),
+            instance=None,
+        )
+        self._lhs = lhs
+        self._rhs = rhs
+
+    def resolve(self, state: "BaseState") -> Union[LHS, RHS]:
+        lhs = resolve_value(self._lhs, state)
+        if lhs:
+            return lhs
+
+        return resolve_value(self._rhs, state)

vellum/workflows/graph/__init__.py

@@ -0,0 +1,3 @@
+from .graph import Graph, GraphTarget
+
+__all__ = ["Graph", "GraphTarget"]

vellum/workflows/graph/graph.py

@@ -0,0 +1,131 @@
+from typing import TYPE_CHECKING, Iterator, List, Set, Type, Union
+
+from orderly_set import OrderedSet
+
+from vellum.workflows.edges.edge import Edge
+from vellum.workflows.types.generics import NodeType
+
+if TYPE_CHECKING:
+    from vellum.workflows.nodes.bases.base import BaseNode
+    from vellum.workflows.ports.port import Port
+
+GraphTargetOfSets = Union[
+    Set[NodeType],
+    Set["Graph"],
+    Set[Union[Type["BaseNode"], "Graph"]],
+]
+
+GraphTarget = Union[
+    Type["BaseNode"],
+    "Graph",
+    GraphTargetOfSets,
+]
+
+
+class Graph:
+    _entrypoints: Set["Port"]
+    _edges: List[Edge]
+    _terminals: Set["Port"]
+
+    def __init__(self, entrypoints: Set["Port"], edges: List[Edge], terminals: Set["Port"]):
+        self._edges = edges
+        self._entrypoints = entrypoints
+        self._terminals = terminals
+
+    @staticmethod
+    def from_port(port: "Port") -> "Graph":
+        ports = {port}
+        return Graph(entrypoints=ports, edges=[], terminals=ports)
+
+    @staticmethod
+    def from_node(node: Type["BaseNode"]) -> "Graph":
+        ports = {port for port in node.Ports}
+        return Graph(entrypoints=ports, edges=[], terminals=ports)
+
+    @staticmethod
+    def from_set(targets: GraphTargetOfSets) -> "Graph":
+        entrypoints = set()
+        edges = OrderedSet[Edge]()
+        terminals = set()
+
+        for target in targets:
+            if isinstance(target, Graph):
+                entrypoints.update(target._entrypoints)
+                edges.update(target._edges)
+                terminals.update(target._terminals)
+            else:
+                entrypoints.update({port for port in target.Ports})
+                terminals.update({port for port in target.Ports})
+
+        return Graph(entrypoints=entrypoints, edges=list(edges), terminals=terminals)
+
+    @staticmethod
+    def from_edge(edge: Edge) -> "Graph":
+        return Graph(entrypoints={edge.from_port}, edges=[edge], terminals={port for port in edge.to_node.Ports})
+
+    def __rshift__(self, other: GraphTarget) -> "Graph":
+        if not self._edges and not self._entrypoints:
+            raise ValueError("Graph instance can only create new edges from nodes within existing edges")
+
+        if isinstance(other, set):
+            new_terminals = set()
+            for elem in other:
+                for final_output_node in self._terminals:
+                    if isinstance(elem, Graph):
+                        midgraph = final_output_node >> set(elem.entrypoints)
+                        self._extend_edges(midgraph.edges)
+                        self._extend_edges(elem.edges)
+                        for other_terminal in elem._terminals:
+                            new_terminals.add(other_terminal)
+                    else:
+                        midgraph = final_output_node >> elem
+                        self._extend_edges(midgraph.edges)
+                        for other_terminal in elem.Ports:
+                            new_terminals.add(other_terminal)
+            self._terminals = new_terminals
+            return self
+
+        if isinstance(other, Graph):
+            for final_output_node in self._terminals:
+                midgraph = final_output_node >> set(other.entrypoints)
+                self._extend_edges(midgraph.edges)
+                self._extend_edges(other.edges)
+            self._terminals = other._terminals
+            return self
+
+        for final_output_node in self._terminals:
+            subgraph = final_output_node >> other
+            self._extend_edges(subgraph.edges)
+        self._terminals = {port for port in other.Ports}
+        return self
+
+    @property
+    def entrypoints(self) -> Iterator[Type["BaseNode"]]:
+        return iter(e.node_class for e in self._entrypoints)
+
+    @property
+    def edges(self) -> Iterator[Edge]:
+        return iter(self._edges)
+
+    @property
+    def nodes(self) -> Iterator[Type["BaseNode"]]:
+        nodes = set()
+        if not self._edges:
+            for node in self.entrypoints:
+                if node not in nodes:
+                    nodes.add(node)
+                    yield node
+            return
+
+        for edge in self._edges:
+            if edge.from_port.node_class not in nodes:
+                nodes.add(edge.from_port.node_class)
+                yield edge.from_port.node_class
+            if edge.to_node not in nodes:
+                nodes.add(edge.to_node)
+                yield edge.to_node
+
+    def _extend_edges(self, edges: Iterator[Edge]) -> None:
+        for edge in edges:
+            if edge not in self._edges:
+                self._edges.append(edge)

vellum/workflows/graph/tests/test_graph.py

@@ -0,0 +1,437 @@
+from vellum.workflows.edges.edge import Edge
+from vellum.workflows.graph.graph import Graph
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.ports.port import Port
+
+
+def test_graph__from_node():
+    # GIVEN a node
+    class MyNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the node
+    graph = Graph.from_node(MyNode)
+
+    # THEN the graph has the node as the entrypoint
+    assert set(graph.entrypoints) == {MyNode}
+
+    # AND one node
+    assert len(list(graph.nodes)) == 1
+
+    # AND no edges
+    assert len(list(graph.edges)) == 0
+
+
+def test_graph__from_edge():
+    # GIVEN an edge
+    class SourceNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    edge = Edge(SourceNode.Ports.default, TargetNode)
+
+    # WHEN we create a graph from the edge
+    graph = Graph.from_edge(edge)
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND two nodes
+    assert len(list(graph.nodes)) == 2
+
+    # AND one edge
+    assert len(list(graph.edges)) == 1
+
+
+def test_graph__node_to_node():
+    # GIVEN two nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the source node to the target node
+    graph = SourceNode >> TargetNode
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND two nodes
+    assert len(list(graph.nodes)) == 2
+
+    # AND one edge
+    assert len(list(graph.edges)) == 1
+
+
+def test_graph__port_to_node():
+    # GIVEN a two nodes, where the source node has a port
+    class SourceNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            output = Port.on_else()
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the source port to the target node
+    graph = SourceNode.Ports.output >> TargetNode
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND two nodes
+    assert len(list(graph.nodes)) == 2
+
+    # AND one edge
+    assert len(list(graph.edges)) == 1
+
+
+def test_graph__graph_to_node():
+    # GIVEN three nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class MiddleNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the source node to the target node via the middle node
+    graph = SourceNode >> MiddleNode >> TargetNode
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2
+
+
+def test_graph__edgeless_graph_to_node():
+    # GIVEN two nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # AND a graph of the first node
+    subgraph = Graph.from_node(SourceNode)
+
+    # WHEN we create a graph from the source graph to the target node
+    graph = subgraph >> TargetNode
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND two nodes
+    assert len(list(graph.nodes)) == 2
+
+    # AND one edge
+    assert len(list(graph.edges)) == 1
+
+    # AND the first node's port has reference to the edge
+    assert set(SourceNode.Ports.default.edges) == set(graph.edges)
+
+
+def test_graph__edgeless_graph_to_graph():
+    # GIVEN two nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # AND a graph of the first node
+    subgraph = Graph.from_node(SourceNode)
+
+    # AND a graph of the second node
+    target_subgraph = Graph.from_node(TargetNode)
+
+    # WHEN we create a graph from the source graph to the target node
+    graph = subgraph >> target_subgraph
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND two nodes
+    assert len(list(graph.nodes)) == 2
+
+    # AND one edge
+    assert len(list(graph.edges)) == 1
+
+    # AND the first node's port has reference to the edges
+    assert set(SourceNode.Ports.default.edges) == set(graph.edges)
+
+
+def test_graph__graph_to_edgeless_graph():
+    # GIVEN three nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class MiddleNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # AND a graph of the first two nodes
+    subgraph = SourceNode >> MiddleNode
+
+    # AND a graph of the third node
+    target_subgraph = Graph.from_node(TargetNode)
+
+    # WHEN we create a graph from the source graph to the target node
+    graph = subgraph >> target_subgraph
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2
+
+
+def test_graph__node_to_graph():
+    # GIVEN three nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class MiddleNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # AND a graph of the last two nodes
+    target_subgraph = MiddleNode >> TargetNode
+
+    # WHEN we create a graph from the source node to the target graph
+    graph = SourceNode >> target_subgraph
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2
+
+
+def test_graph__repeated_edge():
+    # GIVEN two nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the source node to the target node twice
+    graph = SourceNode >> TargetNode >> SourceNode >> TargetNode
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND two nodes
+    assert len(list(graph.nodes)) == 2
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2
+
+    # AND the first node's port has reference to just one edge
+    assert len(list(SourceNode.Ports.default.edges)) == 1
+
+
+def test_graph__node_to_set():
+    # GIVEN three nodes, one with ports
+    class SourceNode(BaseNode):
+        pass
+
+    class MiddleNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            top = Port.on_if(SourceNode.Execution.count.less_than(1))
+            bottom = Port.on_else()
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph with a set
+    graph = SourceNode >> {
+        MiddleNode.Ports.top >> SourceNode,
+        MiddleNode.Ports.bottom >> TargetNode,
+    }
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 3
+
+
+def test_graph__graph_to_set():
+    # GIVEN four nodes, one with ports
+    class SourceNode(BaseNode):
+        pass
+
+    class SecondNode(BaseNode):
+        pass
+
+    class MiddleNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            top = Port.on_if(SourceNode.Execution.count.less_than(1))
+            bottom = Port.on_else()
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph between a graph and a set
+    graph = (
+        SourceNode
+        >> SecondNode
+        >> {
+            MiddleNode.Ports.top >> SourceNode,
+            MiddleNode.Ports.bottom >> TargetNode,
+        }
+    )
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 4
+
+    # AND two edges
+    assert len(list(graph.edges)) == 4
+
+
+def test_graph__graph_set_to_set():
+    # GIVEN five nodes
+    class SourceNode(BaseNode):
+        pass
+
+    class TopNode(BaseNode):
+        pass
+
+    class BottomNode(BaseNode):
+        pass
+
+    class ConditionalNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            top = Port.on_if(SourceNode.Execution.count.less_than(1))
+            bottom = Port.on_else()
+
+    class EndNode(BaseNode):
+        pass
+
+    # WHEN we create a graph that draws an edge between sets
+    graph = (
+        SourceNode
+        >> {
+            TopNode,
+            BottomNode,
+        }
+        >> {
+            ConditionalNode.Ports.top >> SourceNode,
+            ConditionalNode.Ports.bottom >> EndNode,
+        }
+    )
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 5
+
+    # AND two edges
+    assert len(list(graph.edges)) == 6
+
+
+def test_graph__port_to_graph():
+    # GIVEN three nodes where the first node has a port
+    class SourceNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            custom = Port.on_else()
+
+    class MiddleNode(BaseNode):
+        pass
+
+    class EndNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the port to a subgraph
+    graph = SourceNode.Ports.custom >> (MiddleNode >> EndNode)
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2
+
+
+def test_graph__port_to_set():
+    # GIVEN three nodes where the first node has a port
+    class SourceNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            custom = Port.on_else()
+
+    class TopNode(BaseNode):
+        pass
+
+    class BottomNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from the port to the set
+    graph = SourceNode.Ports.custom >> {
+        TopNode,
+        BottomNode,
+    }
+
+    # THEN the graph has the source node as the entrypoint
+    assert set(graph.entrypoints) == {SourceNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2
+
+
+def test_graph__set_to_node():
+    # GIVEN three nodes
+    class TopNode(BaseNode):
+        pass
+
+    class BottomNode(BaseNode):
+        pass
+
+    class TargetNode(BaseNode):
+        pass
+
+    # WHEN we create a graph from a set to a node
+    graph = {
+        TopNode,
+        BottomNode,
+    } >> TargetNode
+
+    # THEN the graph has both the top node and the bottom node as the entrypoints
+    assert set(graph.entrypoints) == {TopNode, BottomNode}
+
+    # AND three nodes
+    assert len(list(graph.nodes)) == 3
+
+    # AND two edges
+    assert len(list(graph.edges)) == 2

vellum/workflows/inputs/__init__.py

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

vellum/workflows/inputs/base.py

@@ -0,0 +1,55 @@
+from typing import Any, Iterator, Tuple, Type
+from typing_extensions import dataclass_transform
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+from vellum.workflows.references import ExternalInputReference, WorkflowInputReference
+from vellum.workflows.references.input import InputReference
+from vellum.workflows.types.utils import get_class_attr_names, infer_types
+
+
+@dataclass_transform(kw_only_default=True)
+class _BaseInputsMeta(type):
+    def __getattribute__(cls, name: str) -> Any:
+        if not name.startswith("_") and name in cls.__annotations__ and issubclass(cls, BaseInputs):
+            instance = vars(cls).get(name)
+            types = infer_types(cls, name)
+
+            if getattr(cls, "__descriptor_class__", None) is ExternalInputReference:
+                return ExternalInputReference(name=name, types=types, instance=instance, inputs_class=cls)
+            else:
+                return WorkflowInputReference(name=name, types=types, instance=instance, inputs_class=cls)
+
+        return super().__getattribute__(name)
+
+    def __iter__(cls) -> Iterator[InputReference]:
+        # We iterate through the inheritance hierarchy to find all the WorkflowInputReference attached to this
+        # Inputs class. __mro__ is the method resolution order, which is the order in which base classes are resolved.
+        for resolved_cls in cls.__mro__:
+            attr_names = get_class_attr_names(resolved_cls)
+            for attr_name in attr_names:
+                attr_value = getattr(resolved_cls, attr_name)
+                if not isinstance(attr_value, (WorkflowInputReference, ExternalInputReference)):
+                    continue
+
+                yield attr_value
+
+
+class BaseInputs(metaclass=_BaseInputsMeta):
+    __parent_class__: Type = type(None)
+
+    def __init__(self, **kwargs: Any) -> None:
+        for name, value in kwargs.items():
+            setattr(self, name, value)
+
+    def __iter__(self) -> Iterator[Tuple[InputReference, Any]]:
+        for input_descriptor in self.__class__:
+            if hasattr(self, input_descriptor.name):
+                yield (input_descriptor, getattr(self, input_descriptor.name))
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.is_instance_schema(cls)