vellum-ai 1.1.4__tar.gz → 1.2.0__tar.gz

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.1.4
+Version: 1.2.0
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -175,7 +175,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
             adornment: JsonObject = {
                 "id": str(node_id),
-                "label": node.__qualname__,
+                "label": self.label,
                 "base": self.get_base().dict(),
                 "attributes": attributes,
             }
@@ -205,7 +205,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         return {
             "id": str(node_id),
-            "label": node.__qualname__,
+            "label": self.label,
             "type": "GENERIC",
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
@@ -417,8 +417,8 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         if explicit_value and explicit_value.comment and docstring:
             comment = (
                 NodeDisplayComment(value=docstring, expanded=explicit_value.comment.expanded)
-                if explicit_value.comment.expanded
-                else NodeDisplayComment(value=docstring)
+                if explicit_value.comment.expanded is not None
+                else NodeDisplayComment(value=docstring, expanded=True)
             )
             return NodeDisplayData(
                 position=explicit_value.position,
@@ -432,7 +432,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         if docstring:
             return NodeDisplayData(
-                comment=NodeDisplayComment(value=docstring),
+                comment=NodeDisplayComment(value=docstring, expanded=True),
             )
 
         return NodeDisplayData()

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -98,5 +98,5 @@ def test_serialize_display_data():
     # THEN the condition should be serialized correctly
     assert data["display_data"] == {
         "position": {"x": 0.0, "y": 0.0},
-        "comment": {"value": "I hope this works"},
+        "comment": {"expanded": True, "value": "I hope this works"},
     }

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -100,6 +100,10 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "trigger": {
+                "id": str(self.get_target_handle_id()),
+                "merge_behavior": node.Trigger.merge_behavior.value,
+            },
             "outputs": [
                 {"id": str(json_display.id), "name": "json", "type": "JSON", "value": None},
                 {"id": str(output_display.id), "name": "text", "type": "STRING", "value": None},

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/vellum/retry_node.py

@@ -39,7 +39,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
 
         adornment: JsonObject = {
             "id": str(node_id),
-            "label": node.__qualname__,
+            "label": self.label,
             "base": self.get_base().dict(),
             "attributes": attributes,
         }

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_retry_node.py

@@ -29,7 +29,7 @@ def test_retry_node_parameters():
     )
 
     retry_adornment = next(
-        adornment for adornment in serialized_node["adornments"] if adornment["label"] == "RetryNode"
+        adornment for adornment in serialized_node["adornments"] if adornment["label"] == "Retry Node"
     )
 
     max_attempts_attribute = next(attr for attr in retry_adornment["attributes"] if attr["name"] == "max_attempts")

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -1,6 +1,7 @@
 from vellum.client.types.prompt_parameters import PromptParameters
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
 from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
@@ -243,7 +244,7 @@ def test_serialize_tool_router_node():
         },
         "display_data": {"position": {"x": 0.0, "y": 0.0}},
         "id": "690c66e1-1e18-4984-b695-84beb0157541",
-        "label": "RouterNode",
+        "label": "Router Node",
         "outputs": [],
         "ports": [
             {
@@ -332,3 +333,76 @@ def test_serialize_tool_router_node():
         "trigger": {"id": "73a96f44-c2dd-40cc-96f6-49b9f914b166", "merge_behavior": "AWAIT_ATTRIBUTES"},
         "type": "GENERIC",
     }
+
+
+def test_serialize_node__tool_calling_node__subworkflow_with_parent_input_reference():
+    """
+    Test that a tool calling node with a subworkflow that references parent inputs serializes correctly
+    """
+
+    # GIVEN a workflow inputs class
+    class MyInputs(BaseInputs):
+        text: str
+        greeting: str
+
+    # AND a code execution node that references parent inputs
+    class CodeExecution(CodeExecutionNode[BaseState, str]):
+        code = ""
+        code_inputs = {
+            "text": MyInputs.text,
+        }
+        runtime = "PYTHON_3_11_6"
+        packages = []
+
+    # AND a subworkflow that uses the code execution node
+    class FunctionSubworkflow(BaseWorkflow):
+        graph = CodeExecution
+
+        class Outputs(BaseWorkflow.Outputs):
+            output = CodeExecution.Outputs.result
+
+    # AND a tool calling node that uses the subworkflow
+    class MyToolCallingNode(ToolCallingNode):
+        ml_model = "gpt-4.1"
+        prompt_inputs = {
+            "text": MyInputs.text,
+        }
+        blocks = []
+        parameters = PromptParameters()
+        functions = [FunctionSubworkflow]
+
+    # AND a workflow with the tool calling node
+    class Workflow(BaseWorkflow[MyInputs, BaseState]):
+        graph = MyToolCallingNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the functions with parent input references
+    my_tool_calling_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(MyToolCallingNode.__id__)
+    )
+
+    functions_attribute = next(
+        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
+    )
+
+    data = functions_attribute["value"]["value"]["value"][0]
+    nodes = data["exec_config"]["workflow_raw_data"]["nodes"]
+    code_exec_node = next((node for node in nodes if node["type"] == "CODE_EXECUTION"), None)
+
+    assert code_exec_node is not None
+    text_input = next((input for input in code_exec_node["inputs"] if input["key"] == "text"), None)
+    assert text_input == {
+        "id": "da92a1c4-d37c-4008-a1ab-c0bcc0cd20d0",
+        "key": "text",
+        "value": {
+            "rules": [
+                {"type": "INPUT_VARIABLE", "data": {"input_variable_id": "6f0c1889-3f08-4c5c-bb24-f7b94169105c"}}
+            ],
+            "combinator": "OR",
+        },
+    }

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -43,7 +43,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
 
         adornment: JsonObject = {
             "id": str(node_id),
-            "label": node.__qualname__,
+            "label": self.label,
             "base": self.get_base().dict(),
             "attributes": attributes,
         }

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py

@@ -47,7 +47,7 @@ def test_serialize_node__retry(serialize_node):
     assert not DeepDiff(
         {
             "id": "188b50aa-e518-4b7b-a5e0-e2585fb1d7b5",
-            "label": "test_serialize_node__retry.<locals>.InnerRetryGenericNode",
+            "label": "Inner Retry Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -68,7 +68,7 @@ def test_serialize_node__retry(serialize_node):
             "adornments": [
                 {
                     "id": "5be7d260-74f7-4734-b31b-a46a94539586",
-                    "label": "RetryNode",
+                    "label": "Retry Node",
                     "base": {
                         "name": "RetryNode",
                         "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
@@ -154,7 +154,7 @@ def test_serialize_node__try(serialize_node):
     assert not DeepDiff(
         {
             "id": str(InnerTryGenericNode.__wrapped_node__.__id__),
-            "label": "test_serialize_node__try.<locals>.InnerTryGenericNode",
+            "label": "Inner Try Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -175,7 +175,7 @@ def test_serialize_node__try(serialize_node):
             "adornments": [
                 {
                     "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
-                    "label": "TryNode",
+                    "label": "Try Node",
                     "base": {
                         "name": "TryNode",
                         "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"],
@@ -248,7 +248,7 @@ def test_serialize_node__stacked():
     assert not DeepDiff(
         {
             "id": "074833b0-e142-4bbc-8dec-209a35e178a3",
-            "label": "test_serialize_node__stacked.<locals>.InnerStackedGenericNode",
+            "label": "Inner Stacked Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -272,7 +272,7 @@ def test_serialize_node__stacked():
             "adornments": [
                 {
                     "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
-                    "label": "TryNode",
+                    "label": "Try Node",
                     "base": {
                         "name": "TryNode",
                         "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"],
@@ -287,7 +287,7 @@ def test_serialize_node__stacked():
                 },
                 {
                     "id": "5be7d260-74f7-4734-b31b-a46a94539586",
-                    "label": "RetryNode",
+                    "label": "Retry Node",
                     "base": {
                         "name": "RetryNode",
                         "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
@@ -351,5 +351,5 @@ def test_serialize_node__adornment_order_matches_decorator_order():
 
     adornments = cast(List[Dict[str, Any]], my_node["adornments"])
     assert len(adornments) == 2
-    assert adornments[0]["label"] == "TryNode"
-    assert adornments[1]["label"] == "RetryNode"
+    assert adornments[0]["label"] == "Try Node"
+    assert adornments[1]["label"] == "Retry Node"

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -34,7 +34,7 @@ def test_serialize_node__constant_value(serialize_node):
     assert not DeepDiff(
         {
             "id": "67e07859-7f67-4287-9854-06ab4199e576",
-            "label": "test_serialize_node__constant_value.<locals>.ConstantValueGenericNode",
+            "label": "Constant Value Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -104,7 +104,7 @@ def test_serialize_node__constant_value_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "73643f17-e49e-47d2-bd01-bb9c3eab6ae9",
-            "label": "test_serialize_node__constant_value_reference.<locals>.ConstantValueReferenceGenericNode",
+            "label": "Constant Value Reference Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -146,7 +146,7 @@ def test_serialize_node__lazy_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "3d6bfe3b-263a-40a6-8a05-98288e9559a4",
-            "label": "test_serialize_node__lazy_reference.<locals>.LazyReferenceGenericNode",
+            "label": "Lazy Reference Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -229,7 +229,7 @@ def test_serialize_node__workflow_input(serialize_node):
     assert not DeepDiff(
         {
             "id": "30116483-6f38-40e0-baf2-32de0e14e9a3",
-            "label": "test_serialize_node__workflow_input.<locals>.WorkflowInputGenericNode",
+            "label": "Workflow Input Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -292,7 +292,7 @@ def test_serialize_node__workflow_input_as_nested_chat_history():
     assert not DeepDiff(
         {
             "id": "11be9d37-0069-4695-a317-14a3b6519d4e",
-            "label": "test_serialize_node__workflow_input_as_nested_chat_history.<locals>.GenericNode",
+            "label": "Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -362,7 +362,7 @@ def test_serialize_node__node_output(serialize_node):
     assert not DeepDiff(
         {
             "id": "7210742f-8c3e-4379-9800-8b4b7f5dd7ed",
-            "label": "test_serialize_node__node_output.<locals>.GenericNodeReferencingOutput",
+            "label": "Generic Node Referencing Output",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -412,7 +412,7 @@ def test_serialize_node__vellum_secret(serialize_node):
     assert not DeepDiff(
         {
             "id": "0e75bd8f-882e-4ab7-8348-061319b574f7",
-            "label": "test_serialize_node__vellum_secret.<locals>.VellumSecretGenericNode",
+            "label": "Vellum Secret Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -465,7 +465,7 @@ def test_serialize_node__node_execution(serialize_node):
     assert not DeepDiff(
         {
             "id": "f42dda6b-e856-49bd-b203-46c9dd66c08b",
-            "label": "test_serialize_node__node_execution.<locals>.GenericNodeReferencingExecutions",
+            "label": "Generic Node Referencing Executions",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -561,7 +561,7 @@ def test_serialize_node__coalesce(serialize_node):
     assert not DeepDiff(
         {
             "id": "bb99f326-7d2a-4b5e-95f3-6039114798da",
-            "label": "test_serialize_node__coalesce.<locals>.CoalesceNodeFinal",
+            "label": "Coalesce Node Final",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -670,3 +670,53 @@ def test_serialize_node__pydantic_with_node_output_reference(serialize_node):
     assert any(
         entry["key"] == "node_ref" and entry["value"]["type"] == "NODE_OUTPUT" for entry in attr_value["entries"]
     )
+
+
+def test_serialize_node__comment_expanded_true_when_content_exists(serialize_node):
+    """
+    Tests that node comment serialization sets expanded=True when comment has content.
+    """
+
+    class NodeWithComment(BaseNode):
+        """This is a test comment for the node."""
+
+        pass
+
+    serialized_node = serialize_node(NodeWithComment)
+
+    # WHEN the node is serialized
+    display_data = serialized_node["display_data"]
+
+    # THEN the comment should have expanded=True
+    assert "comment" in display_data
+    assert display_data["comment"]["value"] == "This is a test comment for the node."
+    assert display_data["comment"]["expanded"] is True
+
+
+def test_serialize_node__comment_expanded_preserved_when_explicitly_set(serialize_node):
+    """
+    Tests that explicitly set expanded value is preserved during serialization.
+    """
+    from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
+    from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+
+    class NodeWithExplicitComment(BaseNode):
+        """This is a test comment."""
+
+        pass
+
+    class NodeWithExplicitCommentDisplay(BaseNodeDisplay[NodeWithExplicitComment]):
+        display_data = NodeDisplayData(comment=NodeDisplayComment(value="Custom comment", expanded=False))
+
+    serialized_node = serialize_node(
+        NodeWithExplicitComment,
+        global_node_displays={NodeWithExplicitComment: NodeWithExplicitCommentDisplay()},
+    )
+
+    # WHEN the node is serialized
+    display_data = serialized_node["display_data"]
+
+    # THEN the comment should preserve expanded=False
+    assert "comment" in display_data
+    assert display_data["comment"]["value"] == "This is a test comment."
+    assert display_data["comment"]["expanded"] is False

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py

@@ -23,7 +23,7 @@ def test_serialize_node__annotated_output(serialize_node):
     assert not DeepDiff(
         {
             "id": "e33ddf79-f48c-4057-ba17-d41a3a60ac98",
-            "label": "test_serialize_node__annotated_output.<locals>.AnnotatedOutputGenericNode",
+            "label": "Annotated Output Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -71,7 +71,7 @@ def test_serialize_node__workflow_input(serialize_node):
     assert not DeepDiff(
         {
             "id": "30116483-6f38-40e0-baf2-32de0e14e9a3",
-            "label": "test_serialize_node__workflow_input.<locals>.WorkflowInputGenericNode",
+            "label": "Workflow Input Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -134,7 +134,7 @@ def test_serialize_node__node_output_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "ac067acc-6a6f-44b1-ae84-428e965ce691",
-            "label": "test_serialize_node__node_output_reference.<locals>.GenericNodeReferencingOutput",
+            "label": "Generic Node Referencing Output",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py

@@ -20,11 +20,10 @@ def test_serialize_node__basic(serialize_node):
         pass
 
     serialized_node = serialize_node(BasicGenericNode)
-
     assert not DeepDiff(
         {
             "id": "8d7cbfe4-72ca-4367-a401-8d28723d2f00",
-            "label": "test_serialize_node__basic.<locals>.BasicGenericNode",
+            "label": "Basic Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -70,7 +69,7 @@ def test_serialize_node__if(serialize_node):
     assert not DeepDiff(
         {
             "id": "bba4b15a-dea0-48c9-a79b-4e12e99db00f",
-            "label": "test_serialize_node__if.<locals>.IfGenericNode",
+            "label": "If Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -132,7 +131,7 @@ def test_serialize_node__if_else(serialize_node):
     assert not DeepDiff(
         {
             "id": "25c9c3f1-4014-47ac-90cf-5216de10d05c",
-            "label": "test_serialize_node__if_else.<locals>.IfElseGenericNode",
+            "label": "If Else Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -201,7 +200,7 @@ def test_serialize_node__if_elif_else(serialize_node):
     assert not DeepDiff(
         {
             "id": "7b2b9cfc-12aa-432c-940d-cbe53e71de9c",
-            "label": "test_serialize_node__if_elif_else.<locals>.IfElifElseGenericNode",
+            "label": "If Elif Else Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -300,7 +299,7 @@ def test_serialize_node__node_output_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "ac067acc-6a6f-44b1-ae84-428e965ce691",
-            "label": "test_serialize_node__node_output_reference.<locals>.GenericNodeReferencingOutput",
+            "label": "Generic Node Referencing Output",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -363,7 +362,7 @@ def test_serialize_node__vellum_secret_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "feb4b331-e25f-4a5c-9840-c5575b1efd5c",
-            "label": "test_serialize_node__vellum_secret_reference.<locals>.GenericNodeReferencingSecret",
+            "label": "Generic Node Referencing Secret",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -429,7 +428,7 @@ def test_serialize_node__execution_count_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "0b4fe8a6-6d0c-464e-9372-10110e2b0e13",
-            "label": "test_serialize_node__execution_count_reference.<locals>.GenericNodeReferencingExecutions",
+            "label": "Generic Node Referencing Executions",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -490,7 +489,7 @@ def test_serialize_node__null(serialize_node):
     assert not DeepDiff(
         {
             "id": "1838ce1f-9c07-4fd0-9fd4-2a3a841ea402",
-            "label": "test_serialize_node__null.<locals>.NullGenericNode",
+            "label": "Null Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -548,7 +547,7 @@ def test_serialize_node__between(serialize_node):
     assert not DeepDiff(
         {
             "id": "f2f5a1f2-a12d-4ce0-bfe9-42190ffe5328",
-            "label": "test_serialize_node__between.<locals>.BetweenGenericNode",
+            "label": "Between Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -617,7 +616,7 @@ def test_serialize_node__or(serialize_node):
     assert not DeepDiff(
         {
             "id": "5386abad-3378-4378-b3a8-831b4b77dc23",
-            "label": "test_serialize_node__or.<locals>.OrGenericNode",
+            "label": "Or Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -700,7 +699,7 @@ def test_serialize_node__and_then_or(serialize_node):
     assert not DeepDiff(
         {
             "id": "4d3995b1-437b-48d9-8878-9f57a8b725f1",
-            "label": "test_serialize_node__and_then_or.<locals>.AndThenOrGenericNode",
+            "label": "And Then Or Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -802,7 +801,7 @@ def test_serialize_node__parenthesized_and_then_or(serialize_node):
     assert not DeepDiff(
         {
             "id": "223864c9-0088-4c05-9b7d-e5b1c9ec936d",
-            "label": "test_serialize_node__parenthesized_and_then_or.<locals>.ParenthesizedAndThenOrGenericNode",
+            "label": "Parenthesized And Then Or Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -904,7 +903,7 @@ def test_serialize_node__or_then_and(serialize_node):
     assert not DeepDiff(
         {
             "id": "a946342e-4ede-4e96-8e3d-f396748d9f7c",
-            "label": "test_serialize_node__or_then_and.<locals>.OrThenAndGenericNode",
+            "label": "Or Then And Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -1005,7 +1004,7 @@ def test_serialize_node__parse_json(serialize_node):
     assert not DeepDiff(
         {
             "id": "bfc3f81b-242a-4f43-9e1c-648223d77768",
-            "label": "test_serialize_node__parse_json.<locals>.ParseJsonGenericNode",
+            "label": "Parse Json Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},

{vellum_ai-1.1.4 → vellum_ai-1.2.0}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py

@@ -2,7 +2,11 @@ from deepdiff import DeepDiff
 
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.core import MergeBehavior
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
 class Inputs(BaseInputs):
@@ -17,7 +21,7 @@ def test_serialize_node__basic(serialize_node):
     assert not DeepDiff(
         {
             "id": "8d7cbfe4-72ca-4367-a401-8d28723d2f00",
-            "label": "test_serialize_node__basic.<locals>.BasicGenericNode",
+            "label": "Basic Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -59,7 +63,7 @@ def test_serialize_node__await_any(serialize_node):
     assert not DeepDiff(
         {
             "id": "42e17f0e-8496-415f-9c72-f85250ba6f0b",
-            "label": "test_serialize_node__await_any.<locals>.AwaitAnyGenericNode",
+            "label": "Await Any Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -101,7 +105,7 @@ def test_serialize_node__await_all(serialize_node):
     assert not DeepDiff(
         {
             "id": "b3e1145a-5f41-456b-9382-6d0a1e828c2f",
-            "label": "test_serialize_node__await_all.<locals>.AwaitAllGenericNode",
+            "label": "Await All Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -132,3 +136,54 @@ def test_serialize_node__await_all(serialize_node):
         serialized_node,
         ignore_order=True,
     )
+
+
+def test_serialize_node__inline_prompt_await_all():
+    """
+    Tests that InlinePromptNode with AWAIT_ALL merge behavior can be defined and serializes without errors.
+    """
+
+    # GIVEN an InlinePromptNode with AWAIT_ALL merge behavior
+    class AwaitAllInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+
+        class Trigger(InlinePromptNode.Trigger):
+            merge_behavior = MergeBehavior.AWAIT_ALL
+
+    class TestWorkflow(BaseWorkflow[Inputs, BaseState]):
+        graph = AwaitAllInlinePromptNode
+
+    # WHEN we serialize the workflow containing the node
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+    serialized = workflow_display.serialize()
+
+    # THEN the workflow should serialize successfully
+    assert "workflow_raw_data" in serialized  # type: ignore
+    assert "nodes" in serialized["workflow_raw_data"]  # type: ignore
+
+    # AND the workflow should contain the InlinePromptNode
+    nodes = serialized["workflow_raw_data"]["nodes"]  # type: ignore
+    prompt_nodes = [node for node in nodes if node["type"] == "PROMPT"]  # type: ignore
+    assert len(prompt_nodes) == 1
+
+    prompt_node = prompt_nodes[0]
+
+    # AND the node should have the correct type and base
+    assert prompt_node["type"] == "PROMPT"  # type: ignore
+    assert prompt_node["base"]["name"] == "InlinePromptNode"  # type: ignore
+    assert prompt_node["base"]["module"] == [  # type: ignore
+        "vellum",
+        "workflows",
+        "nodes",
+        "displayable",
+        "inline_prompt_node",
+        "node",
+    ]
+
+    # AND the node should have the expected structure (InlinePromptNode doesn't serialize trigger info)
+    assert "data" in prompt_node  # type: ignore
+    assert "ml_model_name" in prompt_node["data"]  # type: ignore
+    assert prompt_node["data"]["ml_model_name"] == "gpt-4o"  # type: ignore
+
+    assert prompt_node["trigger"]["merge_behavior"] == "AWAIT_ALL"  # type: ignore