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

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

@@ -8,6 +8,7 @@ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeV
 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_ee.workflows.display.vellum import WorkspaceSecretPointer
 
 _APINodeType = TypeVar("_APINodeType", bound=APINode)
 
@@ -79,16 +80,13 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
             if authorization_type
             else None
         )
-        bearer_token_value_node_input = (
-            create_node_input(
-                node_id=node_id,
-                input_name="bearer_token_value",
-                value=bearer_token_value,
-                display_context=display_context,
-                input_id=self.bearer_token_value_input_id,
-            )
-            if bearer_token_value
-            else None
+        bearer_token_value_node_input = create_node_input(
+            node_id=node_id,
+            input_name="bearer_token_value",
+            value=bearer_token_value,
+            display_context=display_context,
+            input_id=self.bearer_token_value_input_id,
+            pointer_type=WorkspaceSecretPointer,
         )
         api_key_header_key_node_input = (
             create_node_input(
@@ -101,52 +99,52 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
             if api_key_header_key
             else None
         )
-        api_key_header_value_node_input = (
-            create_node_input(
-                node_id=node_id,
-                input_name="api_key_header_value",
-                value=node.api_key_header_value,
-                display_context=display_context,
-                input_id=self.api_key_header_value_input_id,
-            )
-            if headers and api_key_header_key and api_key_header_value
-            else None
+        api_key_header_value_node_input = create_node_input(
+            node_id=node_id,
+            input_name="api_key_header_value",
+            value=api_key_header_value,
+            display_context=display_context,
+            input_id=self.api_key_header_value_input_id,
+            pointer_type=WorkspaceSecretPointer,
         )
 
-        # TODO: Add stable IDs for additional headers
-        # https://app.shortcut.com/vellum/story/5083
-        additional_headers: JsonArray = (
-            [
-                {
-                    "header_key_input_id": create_node_input(
-                        node_id=node_id,
-                        input_name="additional_header_key",
-                        value=key,
-                        display_context=display_context,
-                        input_id=(
-                            self.additional_header_key_input_ids.get(key)
-                            if self.additional_header_key_input_ids
-                            else None
-                        ),
-                    ).id,
-                    "header_value_input_id": create_node_input(
-                        node_id=node_id,
-                        input_name="additional_header_value",
-                        value=value,
-                        display_context=display_context,
-                        input_id=(
-                            self.additional_header_value_input_ids.get(key)
-                            if self.additional_header_value_input_ids
-                            else None
-                        ),
-                    ).id,
-                }
-                for key, value in headers.items()
-                if key not in {api_key_header_key, "Authorization"}
-            ]
-            if headers
-            else []
-        )
+        additional_header_inputs = []
+
+        additional_headers: JsonArray = []
+        if headers:
+            for key, value in headers.items():
+                if key in {api_key_header_key, "Authorization"}:
+                    continue
+
+                header_key_input = create_node_input(
+                    node_id=node_id,
+                    input_name="additional_header_key",
+                    value=key,
+                    display_context=display_context,
+                    input_id=(
+                        self.additional_header_key_input_ids.get(key) if self.additional_header_key_input_ids else None
+                    ),
+                )
+                header_value_input = create_node_input(
+                    node_id=node_id,
+                    input_name="additional_header_value",
+                    value=value,
+                    display_context=display_context,
+                    input_id=(
+                        self.additional_header_value_input_ids.get(key)
+                        if self.additional_header_value_input_ids
+                        else None
+                    ),
+                )
+
+                additional_header_inputs.extend([header_key_input, header_value_input])
+
+                additional_headers.append(
+                    {
+                        "header_key_input_id": header_key_input.id,
+                        "header_value_input_id": header_value_input.id,
+                    }
+                )
 
         inputs = [
             input
@@ -161,6 +159,7 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
             ]
             if input is not None
         ]
+        inputs.extend(additional_header_inputs)
 
         _, text_output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.text)]
         _, json_output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.json)]

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/vellum.py

@@ -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

@@ -14,7 +14,6 @@ from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
 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,12 +56,12 @@ 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(
                 {
@@ -270,12 +269,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,