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

vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py

@@ -0,0 +1,58 @@
+from uuid import UUID
+from typing import Any, Generic, Optional, TypeVar
+
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum.workflows.nodes import SubworkflowDeploymentNode
+from vellum.workflows.types.core import JsonObject
+from vellum.workflows.vellum_client import create_vellum_client
+
+_SubworkflowDeploymentNodeType = TypeVar("_SubworkflowDeploymentNodeType", bound=SubworkflowDeploymentNode)
+
+
+class BaseSubworkflowDeploymentNodeDisplay(
+    BaseNodeVellumDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
+):
+    def serialize(
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+    ) -> JsonObject:
+        node = self._node
+        node_id = self.node_id
+
+        subworkflow_inputs = raise_if_descriptor(node.subworkflow_inputs)
+        node_inputs = [
+            create_node_input(
+                node_id=node_id,
+                input_name=variable_name,
+                value=variable_value,
+                display_context=display_context,
+                input_id=None,
+            )
+            for variable_name, variable_value in subworkflow_inputs.items()
+        ]
+
+        # TODO: Pass through the name instead of retrieving the ID
+        # https://app.shortcut.com/vellum/story/4702
+        vellum_client = create_vellum_client()
+        deployment = vellum_client.workflow_deployments.retrieve(
+            id=str(raise_if_descriptor(node.deployment)),
+        )
+
+        return {
+            "id": str(node_id),
+            "type": "SUBWORKFLOW",
+            "inputs": [node_input.dict() for node_input in node_inputs],
+            "data": {
+                "label": self.label,
+                "error_output_id": str(error_output_id) if error_output_id else None,
+                "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
+                "target_handle_id": str(self.get_target_handle_id()),
+                "variant": "DEPLOYMENT",
+                "workflow_deployment_id": str(deployment.id),
+                "release_tag": raise_if_descriptor(node.release_tag),
+            },
+            "display_data": self.get_display_data().dict(),
+            "definition": self.get_definition().dict(),
+        }

vellum_ee/workflows/display/nodes/vellum/templating_node.py

@@ -0,0 +1,67 @@
+from uuid import UUID
+from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
+
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum.workflows.nodes.core.templating_node import TemplatingNode
+from vellum.workflows.types.core import JsonObject
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+
+_TemplatingNodeType = TypeVar("_TemplatingNodeType", bound=TemplatingNode)
+
+
+class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
+    output_id: ClassVar[Optional[UUID]] = None
+    template_input_id: ClassVar[Optional[UUID]] = None
+    input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+
+    def serialize(
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+    ) -> JsonObject:
+        node = self._node
+        node_id = self.node_id
+
+        template_node_input = create_node_input(
+            node_id=node_id,
+            input_name="template",
+            value=node.template,
+            display_context=display_context,
+            input_id=self.template_input_id,
+        )
+        template_node_inputs = raise_if_descriptor(node.inputs)
+        template_inputs = [
+            create_node_input(
+                node_id=node_id,
+                input_name=variable_name,
+                value=variable_value,
+                display_context=display_context,
+                input_id=self.input_ids_by_name.get("template"),
+            )
+            for variable_name, variable_value in template_node_inputs.items()
+        ]
+        node_inputs = [template_node_input, *template_inputs]
+
+        # Misc type ignore is due to `node.Outputs` being generic
+        # https://app.shortcut.com/vellum/story/4784
+        output_descriptor = node.Outputs.result  # type: ignore [misc]
+        _, output_display = display_context.node_output_displays[output_descriptor]
+        inferred_output_type = primitive_type_to_vellum_variable_type(output_descriptor)
+
+        return {
+            "id": str(node_id),
+            "type": "TEMPLATING",
+            "inputs": [node_input.dict() for node_input in node_inputs],
+            "data": {
+                "label": self.label,
+                "output_id": str(output_display.id),
+                "error_output_id": str(error_output_id) if error_output_id else None,
+                "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
+                "target_handle_id": str(self.get_target_handle_id()),
+                "template_node_input_id": str(template_node_input.id),
+                "output_type": inferred_output_type,
+            },
+            "display_data": self.get_display_data().dict(),
+            "definition": self.get_definition().dict(),
+        }

vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -0,0 +1,106 @@
+import pytest
+from uuid import UUID, uuid4
+from typing import List, cast
+
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.vellum import (
+    ConstantValuePointer,
+    InputVariableData,
+    InputVariablePointer,
+    NodeDisplayData,
+    NodeInputValuePointerRule,
+    NodeOutputData,
+    NodeOutputPointer,
+    StringVellumValue,
+    WorkflowInputsVellumDisplayOverrides,
+    WorkflowMetaVellumDisplay,
+)
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.references import OutputReference, WorkflowInputReference
+
+
+class Inputs(BaseInputs):
+    example_workflow_input: str
+
+
+class MyNodeA(BaseNode):
+    example_node_input = Inputs.example_workflow_input
+
+    class Outputs(BaseOutputs):
+        output: str
+
+
+class MyNodeB(BaseNode):
+    example = MyNodeA.Outputs.output
+    fallback_example = MyNodeA.Outputs.output.coalesce(Inputs.example_workflow_input).coalesce("fallback")
+
+
+@pytest.mark.parametrize(
+    ["descriptor", "expected_rules"],
+    [
+        (
+            MyNodeB.example,
+            [
+                NodeOutputPointer(
+                    data=NodeOutputData(
+                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
+                    )
+                )
+            ],
+        ),
+        (
+            MyNodeB.fallback_example,
+            [
+                NodeOutputPointer(
+                    type="NODE_OUTPUT",
+                    data=NodeOutputData(
+                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
+                    ),
+                ),
+                InputVariablePointer(
+                    type="INPUT_VARIABLE",
+                    data=InputVariableData(input_variable_id="a154c29d-fac0-4cd0-ba88-bc52034f5470"),
+                ),
+                ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(type="STRING", value="fallback")),
+            ],
+        ),
+    ],
+)
+def test_create_node_input_value_pointer_rules(
+    descriptor: BaseDescriptor, expected_rules: List[NodeInputValuePointerRule]
+) -> None:
+    rules = create_node_input_value_pointer_rules(
+        descriptor,
+        WorkflowDisplayContext(
+            workflow_display_class=VellumWorkflowDisplay,
+            workflow_display=WorkflowMetaVellumDisplay(
+                entrypoint_node_id=uuid4(),
+                entrypoint_node_source_handle_id=uuid4(),
+                entrypoint_node_display=NodeDisplayData(),
+            ),
+            workflow_input_displays={
+                cast(WorkflowInputReference, Inputs.example_workflow_input): WorkflowInputsVellumDisplayOverrides(
+                    id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
+                ),
+            },
+            node_output_displays={
+                cast(OutputReference, MyNodeA.Outputs.output): (
+                    MyNodeA,
+                    NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
+                ),
+            },
+            node_displays={
+                MyNodeA: BaseNodeVellumDisplay(MyNodeA),
+            },
+        ),
+    )
+    assert rules == expected_rules

vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -0,0 +1,38 @@
+from uuid import UUID
+from typing import Any, ClassVar, Generic, Optional, TypeVar
+
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum.workflows.nodes.core.try_node.node import TryNode
+from vellum.workflows.nodes.utils import get_wrapped_node
+from vellum.workflows.types.core import JsonObject
+
+_TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
+
+
+class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeType]):
+    error_output_id: ClassVar[Optional[UUID]] = None
+
+    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+        node = self._node
+
+        try:
+            inner_node = get_wrapped_node(node)
+        except TypeError:
+            raise NotImplementedError(
+                "Unable to serialize Try Nodes that wrap subworkflows containing more than one Node."
+            )
+
+        # We need the node display class of the underlying node because
+        # it contains the logic for serializing the node and potential display overrides
+        node_display_class = get_node_display_class(BaseNodeVellumDisplay, inner_node)
+        node_display = node_display_class(inner_node)
+
+        serialized_node = node_display.serialize(
+            display_context,
+            error_output_id=str(self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id")),
+        )
+
+        return serialized_node

vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -0,0 +1,76 @@
+from uuid import UUID
+from typing import Any, List, Optional, cast
+
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule, primitive_to_vellum_value
+from vellum_ee.workflows.display.vellum import (
+    ConstantValuePointer,
+    NodeInput,
+    NodeInputValuePointer,
+    NodeInputValuePointerRule,
+)
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.references import NodeReference, OutputReference
+
+
+def create_node_input(
+    node_id: UUID,
+    input_name: str,
+    value: Any,
+    display_context: WorkflowDisplayContext,
+    input_id: Optional[UUID],
+) -> NodeInput:
+    input_id = input_id or uuid4_from_hash(f"{node_id}|{input_name}")
+    if (
+        isinstance(value, OutputReference)
+        and value.outputs_class._node_class
+        and has_wrapped_node(value.outputs_class._node_class)
+    ):
+        wrapped_node = get_wrapped_node(value.outputs_class._node_class)
+        if wrapped_node._is_wrapped_node:
+            value = getattr(wrapped_node.Outputs, value.name)
+
+    rules = create_node_input_value_pointer_rules(value, display_context)
+    return NodeInput(
+        id=str(input_id),
+        key=input_name,
+        value=NodeInputValuePointer(
+            rules=rules,
+            combinator="OR",
+        ),
+    )
+
+
+def create_node_input_value_pointer_rules(
+    value: Any,
+    display_context: WorkflowDisplayContext,
+    existing_rules: Optional[List[NodeInputValuePointerRule]] = None,
+) -> List[NodeInputValuePointerRule]:
+    node_input_value_pointer_rules: List[NodeInputValuePointerRule] = existing_rules or []
+
+    if isinstance(value, BaseDescriptor):
+        if isinstance(value, NodeReference):
+            if not value.instance:
+                raise ValueError(f"Expected NodeReference {value.name} to have an instance")
+            value = cast(BaseDescriptor, value.instance)
+
+        if isinstance(value, CoalesceExpression):
+            # Recursively handle the left-hand side
+            lhs_rules = create_node_input_value_pointer_rules(value.lhs, display_context, [])
+            node_input_value_pointer_rules.extend(lhs_rules)
+
+            # Handle the right-hand side
+            if not isinstance(value.rhs, CoalesceExpression):
+                rhs_rules = create_node_input_value_pointer_rules(value.rhs, display_context, [])
+                node_input_value_pointer_rules.extend(rhs_rules)
+        else:
+            # Non-CoalesceExpression case
+            node_input_value_pointer_rules.append(create_node_input_value_pointer_rule(value, display_context))
+    else:
+        vellum_variable_value = primitive_to_vellum_value(value)
+        node_input_value_pointer_rules.append(ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_variable_value))
+
+    return node_input_value_pointer_rules