vellum-ai 0.11.0__py3-none-any.whl → 0.11.3__py3-none-any.whl

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, List, Optional, cast
+from typing import Any, List, Optional, Type, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
@@ -13,6 +13,8 @@ from vellum_ee.workflows.display.vellum import (
     NodeInput,
     NodeInputValuePointer,
     NodeInputValuePointerRule,
+    WorkspaceSecretData,
+    WorkspaceSecretPointer,
 )
 
 
@@ -22,6 +24,7 @@ def create_node_input(
     value: Any,
     display_context: WorkflowDisplayContext,
     input_id: Optional[UUID],
+    pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
 ) -> NodeInput:
     input_id = input_id or uuid4_from_hash(f"{node_id}|{input_name}")
     if (
@@ -33,7 +36,7 @@ def create_node_input(
         if wrapped_node._is_wrapped_node:
             value = getattr(wrapped_node.Outputs, value.name)
 
-    rules = create_node_input_value_pointer_rules(value, display_context)
+    rules = create_node_input_value_pointer_rules(value, display_context, pointer_type=pointer_type)
     return NodeInput(
         id=str(input_id),
         key=input_name,
@@ -48,6 +51,7 @@ def create_node_input_value_pointer_rules(
     value: Any,
     display_context: WorkflowDisplayContext,
     existing_rules: Optional[List[NodeInputValuePointerRule]] = None,
+    pointer_type: Optional[Type[NodeInputValuePointerRule]] = None,
 ) -> List[NodeInputValuePointerRule]:
     node_input_value_pointer_rules: List[NodeInputValuePointerRule] = existing_rules or []
 
@@ -59,18 +63,34 @@ def create_node_input_value_pointer_rules(
 
         if isinstance(value, CoalesceExpression):
             # Recursively handle the left-hand side
-            lhs_rules = create_node_input_value_pointer_rules(value.lhs, display_context, [])
+            lhs_rules = create_node_input_value_pointer_rules(value.lhs, display_context, [], pointer_type=pointer_type)
             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, [])
+                rhs_rules = create_node_input_value_pointer_rules(
+                    value.rhs, display_context, [], pointer_type=pointer_type
+                )
                 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))
+        pointer = create_pointer(value, pointer_type)
+        node_input_value_pointer_rules.append(pointer)
 
     return node_input_value_pointer_rules
+
+
+def create_pointer(
+    value: Any,
+    pointer_type: Optional[Type[NodeInputValuePointerRule]] = None,
+) -> NodeInputValuePointerRule:
+    if value is None:
+        if pointer_type is WorkspaceSecretPointer:
+            return WorkspaceSecretPointer(
+                type="WORKSPACE_SECRET", data=WorkspaceSecretData(type="STRING", workspace_secret_id=None)
+            )
+
+    vellum_variable_value = primitive_to_vellum_value(value)
+    return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_variable_value)

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py

@@ -119,6 +119,16 @@ def test_serialize_workflow(vellum_client):
                         "combinator": "OR",
                     },
                 },
+                {
+                    "id": "fee5e3c9-442a-4922-ba80-5ee07361cea7",
+                    "key": "bearer_token_value",
+                    "value": {
+                        "rules": [
+                            {"type": "WORKSPACE_SECRET", "data": {"type": "STRING", "workspace_secret_id": None}}
+                        ],
+                        "combinator": "OR",
+                    },
+                },
                 {
                     "id": "2fcdfbc3-8095-4277-bb4a-a201fd326b54",
                     "key": "api_key_header_key",
@@ -143,6 +153,24 @@ def test_serialize_workflow(vellum_client):
                         "combinator": "OR",
                     },
                 },
+                {
+                    "id": "57c31247-998a-430d-bb62-bf50eca7df35",
+                    "key": "additional_header_key",
+                    "value": {
+                        "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "STRING", "value": "additional_header"}}],
+                        "combinator": "OR",
+                    },
+                },
+                {
+                    "id": "47b32274-f19b-4c15-b788-55c069c311c5",
+                    "key": "additional_header_value",
+                    "value": {
+                        "rules": [
+                            {"type": "CONSTANT_VALUE", "data": {"type": "STRING", "value": "additional header value"}}
+                        ],
+                        "combinator": "OR",
+                    },
+                },
             ],
             "data": {
                 "label": "Simple A P I Node",
@@ -153,7 +181,7 @@ def test_serialize_workflow(vellum_client):
                 "method_input_id": "fd61b5ac-39f9-4cfe-a839-f8ce78c202df",
                 "body_input_id": "c3a17ceb-e201-4025-b18a-9162aac7705e",
                 "authorization_type_input_id": "3092bf23-3202-4f3e-874c-9a33ccc73459",
-                "bearer_token_value_input_id": None,
+                "bearer_token_value_input_id": "fee5e3c9-442a-4922-ba80-5ee07361cea7",
                 "api_key_header_key_input_id": "2fcdfbc3-8095-4277-bb4a-a201fd326b54",
                 "api_key_header_value_input_id": "d794bb51-a419-4fd8-be63-dfaf4166e831",
                 "additional_headers": [

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py

@@ -218,7 +218,7 @@ def test_serialize_workflow__await_all():
                 "source_node_id": "59243c65-053f-4ea6-9157-3f3edb1477bf",
                 "source_handle_id": "b9c5f52b-b714-46e8-a09c-38b4e770dd36",
                 "target_node_id": "37c10e8a-771b-432b-a767-31f5007851f0",
-                "target_handle_id": "0efd256f-f5f6-45fe-9adb-651780f5e63d",
+                "target_handle_id": "42eeb66c-9792-4609-8c71-3a56f668f4dc",
                 "type": "DEFAULT",
             },
             {
@@ -226,7 +226,7 @@ def test_serialize_workflow__await_all():
                 "source_node_id": "127ef456-91bc-43c6-bd8b-1772db5e3cb5",
                 "source_handle_id": "b0bd17f3-4ce6-4232-9666-ec8afa161bf2",
                 "target_node_id": "37c10e8a-771b-432b-a767-31f5007851f0",
-                "target_handle_id": "0efd256f-f5f6-45fe-9adb-651780f5e63d",
+                "target_handle_id": "f40ff7fb-de1b-4aa4-ba3c-7630f7357cbf",
                 "type": "DEFAULT",
             },
             {

vellum_ee/workflows/display/vellum.py

@@ -74,7 +74,9 @@ class WorkflowMetaVellumDisplay(WorkflowMetaVellumDisplayOverrides):
 
 @dataclass
 class WorkflowInputsVellumDisplayOverrides(WorkflowInputsDisplay, WorkflowInputsDisplayOverrides):
-    pass
+    name: Optional[str] = None
+    required: Optional[bool] = None
+    color: Optional[str] = None
 
 
 @dataclass
@@ -115,8 +117,6 @@ class WorkflowOutputVellumDisplayOverrides(WorkflowOutputDisplay, WorkflowOutput
     target_handle_id: UUID
     display_data: NodeDisplayData
     edge_id: UUID
-    # source_node_id: Optional[UUID]
-    # source_handle_id: Optional[UUID]
 
 
 @dataclass
@@ -202,7 +202,7 @@ class InputVariablePointer(UniversalBaseModel):
 
 class WorkspaceSecretData(UniversalBaseModel):
     type: VellumVariableType
-    workspace_secret_id: str
+    workspace_secret_id: Optional[str] = None
 
 
 class WorkspaceSecretPointer(UniversalBaseModel):

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -12,9 +12,9 @@ from vellum.workflows.references import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
+from vellum_ee.workflows.display.nodes import BaseMergeNodeDisplay
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import PortDisplay
-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.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type, primitive_to_vellum_value
@@ -57,17 +57,17 @@ class VellumWorkflowDisplay(
     def serialize(self, raise_errors: bool = True) -> JsonObject:
         input_variables: JsonArray = []
         for workflow_input, workflow_input_display in self.display_context.workflow_input_displays.items():
-            default = (
-                primitive_to_vellum_value(raise_if_descriptor(workflow_input.instance))
-                if workflow_input.instance
-                else None
+            default = primitive_to_vellum_value(workflow_input.instance) if workflow_input.instance else None
+            required = (
+                workflow_input_display.required
+                if workflow_input_display.required is not None
+                else type(None) not in workflow_input.types
             )
-            required = type(None) not in workflow_input.types
 
             input_variables.append(
                 {
                     "id": str(workflow_input_display.id),
-                    "key": workflow_input.name,
+                    "key": workflow_input_display.name or workflow_input.name,
                     "type": infer_vellum_variable_type(workflow_input),
                     "default": default.dict() if default else None,
                     "required": required,
@@ -270,12 +270,16 @@ class VellumWorkflowDisplay(
         self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsVellumDisplayOverrides] = None
     ) -> WorkflowInputsVellumDisplay:
         workflow_input_id: UUID
+        required = None
+        color = None
         if overrides:
             workflow_input_id = overrides.id
+            required = overrides.required
+            color = overrides.color
         else:
             workflow_input_id = uuid4_from_hash(f"{self.workflow_id}|inputs|id|{workflow_input.name}")
 
-        return WorkflowInputsVellumDisplay(id=workflow_input_id)
+        return WorkflowInputsVellumDisplay(id=workflow_input_id, required=required, color=color)
 
     def _generate_entrypoint_display(
         self,
@@ -363,7 +367,12 @@ class VellumWorkflowDisplay(
 
         target_node_display = node_displays[target_node]
         target_node_id = target_node_display.node_id
-        target_handle_id = target_node_display.get_target_handle_id()
+
+        target_handle_id: UUID
+        if isinstance(target_node_display, BaseMergeNodeDisplay):
+            target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
+        else:
+            target_handle_id = target_node_display.get_target_handle_id()
 
         return self._generate_edge_display_from_source(
             source_node_id, source_handle_id, target_node_id, target_handle_id, overrides