vellum-ai 0.14.63__py3-none-any.whl → 1.13.5__py3-none-any.whl

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

@@ -1,14 +1,14 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 from vellum import (
     MetadataFilterConfigRequest,
     VellumValueLogicalConditionGroupRequest,
     VellumValueLogicalConditionRequest,
 )
+from vellum.workflows.nodes.displayable.bases.types import MetadataLogicalCondition, MetadataLogicalConditionGroup
 from vellum.workflows.nodes.displayable.search_node import SearchNode
-from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
@@ -42,16 +42,14 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
     }
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
         node_inputs = self._generate_search_node_inputs(node_id, node, display_context)
 
-        _, results_output_display = display_context.global_node_output_displays[
-            cast(OutputReference, node.Outputs.results)
-        ]
-        _, text_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
+        results_output_display = self.get_node_output_display(node.Outputs.results)
+        text_output_display = self.get_node_output_display(node.Outputs.text)
 
         return {
             "id": str(node_id),
@@ -73,10 +71,7 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
                 "external_id_filters_node_input_id": str(node_inputs["external_id_filters"].id),
                 "metadata_filters_node_input_id": str(node_inputs["metadata_filters"].id),
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
 
     def _generate_search_node_inputs(
@@ -88,7 +83,8 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
         node_inputs: Dict[str, NodeInput] = {}
 
         options = raise_if_descriptor(node.options)
-        filters = options.filters if options else None
+        raw_filters = raise_if_descriptor(node.filters)
+        filters = raw_filters if raw_filters else options.filters if options else None
 
         external_id_filters = filters.external_ids if filters else None
 
@@ -104,17 +100,21 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
                 raw_metadata_filters, display_context=display_context
             )
 
-        result_merging = options.result_merging if options else None
+        raw_result_merging = raise_if_descriptor(node.result_merging)
+        result_merging = raw_result_merging if raw_result_merging else options.result_merging if options else None
         result_merging_enabled = True if result_merging and result_merging.enabled else False
 
         raw_weights = raise_if_descriptor(node.weights)
         weights = raw_weights if raw_weights is not None else options.weights if options is not None else None
 
+        raw_limit = raise_if_descriptor(node.limit)
+        limit = raw_limit if raw_limit is not None else options.limit if options is not None else None
+
         node_input_names_and_values = [
-            ("query", node.query),
-            ("document_index_id", node.document_index),
+            ("query", raise_if_descriptor(node.query)),
+            ("document_index_id", raise_if_descriptor(node.document_index)),
             ("weights", weights.dict() if weights else None),
-            ("limit", options.limit if options else None),
+            ("limit", limit),
             ("separator", raise_if_descriptor(node.chunk_separator)),
             (
                 "result_merging_enabled",
@@ -141,7 +141,12 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
 
     def _serialize_logical_expression(
         self,
-        logical_expression: Union[VellumValueLogicalConditionGroupRequest, VellumValueLogicalConditionRequest],
+        logical_expression: Union[
+            VellumValueLogicalConditionGroupRequest,
+            VellumValueLogicalConditionRequest,
+            MetadataLogicalConditionGroup,
+            MetadataLogicalCondition,
+        ],
         display_context: WorkflowDisplayContext,
         path: List[int] = [],
     ) -> Tuple[JsonObject, List[NodeInput]]:
@@ -175,10 +180,10 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
 
             lhs_query_input_id: UUID = self.metadata_filter_input_id_by_operand_id.get(
                 UUID(lhs_variable_id)
-            ) or uuid4_from_hash(f"{self.node_id}|{hash(tuple(path))}")
+            ) or uuid4_from_hash(f"{self.node_id}|lhs|{hash(tuple(path))}")
             rhs_query_input_id: UUID = self.metadata_filter_input_id_by_operand_id.get(
                 UUID(rhs_variable_id)
-            ) or uuid4_from_hash(f"{self.node_id}|{hash(tuple(path))}")
+            ) or uuid4_from_hash(f"{self.node_id}|rhs|{hash(tuple(path))}")
 
             return (
                 {
@@ -206,5 +211,57 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
                     ),
                 ],
             )
+
+        elif isinstance(logical_expression, MetadataLogicalConditionGroup):
+            conditions = []
+            variables = []
+            for idx, metadata_condition in enumerate(logical_expression.conditions):
+                serialized_condition, serialized_variables = self._serialize_logical_expression(
+                    metadata_condition, display_context=display_context, path=path + [idx]
+                )
+                conditions.append(serialized_condition)
+                variables.extend(serialized_variables)
+
+            return (
+                {
+                    "type": "LOGICAL_CONDITION_GROUP",
+                    "combinator": logical_expression.combinator,
+                    "conditions": conditions,
+                    "negated": logical_expression.negated,
+                },
+                variables,
+            )
+
+        elif isinstance(logical_expression, MetadataLogicalCondition):
+            lhs_variable = logical_expression.lhs_variable
+            rhs_variable = logical_expression.rhs_variable
+
+            lhs_query_input_id = uuid4_from_hash(f"{self.node_id}|lhs|{hash(tuple(path))}")
+            rhs_query_input_id = uuid4_from_hash(f"{self.node_id}|rhs|{hash(tuple(path))}")
+
+            return (
+                {
+                    "type": "LOGICAL_CONDITION",
+                    "lhs_variable_id": str(lhs_query_input_id),
+                    "operator": logical_expression.operator,
+                    "rhs_variable_id": str(rhs_query_input_id),
+                },
+                [
+                    create_node_input(
+                        self.node_id,
+                        f"vellum-query-builder-variable-{lhs_query_input_id}",
+                        lhs_variable,
+                        display_context,
+                        input_id=lhs_query_input_id,
+                    ),
+                    create_node_input(
+                        self.node_id,
+                        f"vellum-query-builder-variable-{rhs_query_input_id}",
+                        rhs_variable,
+                        display_context,
+                        input_id=rhs_query_input_id,
+                    ),
+                ],
+            )
         else:
             raise ValueError(f"Unsupported logical expression type: {type(logical_expression)}")

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

@@ -1,9 +1,13 @@
 from uuid import UUID
 from typing import Generic, Optional, TypeVar
 
+from vellum.client import Vellum as VellumClient
+from vellum.utils.uuid import is_valid_uuid
+from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import SubworkflowDeploymentNode
 from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 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
@@ -16,13 +20,50 @@ class BaseSubworkflowDeploymentNodeDisplay(
 ):
     __serializable_inputs__ = {SubworkflowDeploymentNode.subworkflow_inputs}
 
+    _deployment_id: Optional[str] = None
+    _release_tag: Optional[str] = None
+
+    def build(self, client: VellumClient) -> None:
+        node = self._node
+        deployment_descriptor_id = str(raise_if_descriptor(node.deployment))
+        self._release_tag = raise_if_descriptor(node.release_tag)
+
+        deployment_release = client.workflow_deployments.retrieve_workflow_deployment_release(
+            id=deployment_descriptor_id,
+            release_id_or_release_tag=self._release_tag,
+        )
+        self._deployment_id = str(deployment_release.deployment.id)
+        output_variables_by_key = {var.key: var for var in deployment_release.workflow_version.output_variables}
+
+        for output in node.Outputs:
+            original_output_display = self.output_display.get(output)
+            if original_output_display and not original_output_display._is_implicit:
+                continue
+
+            output_variable = output_variables_by_key.get(output.name)
+            if not output_variable or not is_valid_uuid(output_variable.id):
+                continue
+
+            output_id = UUID(output_variable.id)
+            self.output_display[output] = NodeOutputDisplay(
+                id=output_id,
+                name=output.name,
+            )
+            self._node.__output_ids__[output.name] = output_id
+
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
         subworkflow_inputs = raise_if_descriptor(node.subworkflow_inputs)
+
+        if isinstance(subworkflow_inputs, BaseInputs):
+            input_items = [(input_descriptor.name, input_value) for input_descriptor, input_value in subworkflow_inputs]
+        else:
+            input_items = list(subworkflow_inputs.items())
+
         node_inputs = [
             create_node_input(
                 node_id=node_id,
@@ -34,13 +75,9 @@ class BaseSubworkflowDeploymentNodeDisplay(
                 )
                 or self.node_input_ids_by_name.get(variable_name),
             )
-            for variable_name, variable_value in subworkflow_inputs.items()
+            for variable_name, variable_value in input_items
         ]
 
-        deployment = display_context.client.workflow_deployments.retrieve(
-            id=str(raise_if_descriptor(node.deployment)),
-        )
-
         return {
             "id": str(node_id),
             "type": "SUBWORKFLOW",
@@ -51,11 +88,8 @@ class BaseSubworkflowDeploymentNodeDisplay(
                 "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),
+                "workflow_deployment_id": self._deployment_id,
+                "release_tag": self._release_tag,
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context),
         }

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

@@ -18,7 +18,7 @@ class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_T
     __serializable_inputs__ = {TemplatingNode.inputs}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -50,7 +50,7 @@ class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_T
         # 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.global_node_output_displays[output_descriptor]
+        output_display = self.get_node_output_display(output_descriptor)
         inferred_output_type = primitive_type_to_vellum_variable_type(output_descriptor)
 
         return {
@@ -66,8 +66,5 @@ class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_T
                 "template_node_input_id": str(template_node_input.id),
                 "output_type": inferred_output_type,
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }

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

@@ -0,0 +1,65 @@
+from vellum.workflows.constants import APIRequestMethod
+from vellum.workflows.nodes.displayable.api_node.node import APINode
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_serialize_node__api_node_with_timeout():
+    # GIVEN an API node with timeout specified
+    class MyAPINode(APINode):
+        url = "https://api.example.com"
+        method = APIRequestMethod.GET
+        timeout = 30  # This is the key attribute we're testing
+
+    # AND a workflow with the API node
+    class Workflow(BaseWorkflow):
+        graph = MyAPINode
+
+    # 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 timeout
+    my_api_node = next(node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "API")
+
+    # Assert that timeout is present in the serialized attributes
+    timeout_attribute = next((attr for attr in my_api_node.get("attributes", []) if attr["name"] == "timeout"), None)
+
+    assert timeout_attribute is not None, "timeout attribute should be present in serialized attributes"
+    assert timeout_attribute["value"]["type"] == "CONSTANT_VALUE"
+    assert timeout_attribute["value"]["value"]["type"] == "NUMBER"
+    assert timeout_attribute["value"]["value"]["value"] == 30.0
+
+
+def test_serialize_node__api_node_outputs():
+    """
+    Tests that API node serialization includes the outputs array.
+    """
+
+    class MyAPINode(APINode):
+        url = "https://api.example.com"
+        method = APIRequestMethod.GET
+
+    class Workflow(BaseWorkflow):
+        graph = MyAPINode
+
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    my_api_node = next(node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "API")
+
+    assert "outputs" in my_api_node
+    outputs = my_api_node["outputs"]
+    assert len(outputs) == 4
+
+    text_output = next(output for output in outputs if output["name"] == "text")
+    assert text_output["type"] == "STRING"
+    assert text_output["value"] is None
+
+    json_output = next(output for output in outputs if output["name"] == "json")
+    assert json_output["type"] == "JSON"
+    assert json_output["value"] is None
+
+    status_code_output = next(output for output in outputs if output["name"] == "status_code")
+    assert status_code_output["type"] == "NUMBER"
+    assert status_code_output["value"] is None

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

@@ -3,10 +3,13 @@ from uuid import UUID
 from typing import Type
 
 from vellum.client.core.api_error import ApiError
+from vellum.workflows.environment import EnvironmentVariables
 from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
 from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.state.base import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.vellum.code_execution_node import BaseCodeExecutionNodeDisplay
+from vellum_ee.workflows.display.utils.exceptions import NodeValidationError
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -31,7 +34,7 @@ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[Cod
 @pytest.mark.parametrize(
     ["GetDisplayClass", "expected_input_id"],
     [
-        (_no_display_class, "a5dbe403-0b00-4df6-b8f7-ed5f7794b003"),
+        (_no_display_class, "20ff166f-af59-4515-8ff5-205226c01aa4"),
         (_display_class_with_node_input_ids_by_name, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
         (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
     ],
@@ -80,7 +83,7 @@ def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
             },
         },
         {
-            "id": "9774d864-c76d-4a1a-8181-b632ed3ab87c",
+            "id": "50678b9f-bdea-41c0-bc3d-425ea38466ee",
             "key": "code",
             "value": {
                 "combinator": "OR",
@@ -96,7 +99,7 @@ def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
             },
         },
         {
-            "id": "34742235-5699-45cd-9d34-bce3745e743d",
+            "id": "6d2840f3-a5c1-4376-8616-ced4fffc6cf2",
             "key": "runtime",
             "value": {
                 "combinator": "OR",
@@ -153,3 +156,233 @@ def test_serialize_node__with_unresolved_secret_references(vellum_client):
     # warnings = list(workflow_display.errors)
     # assert len(warnings) == 1
     # assert "Failed to resolve secret reference 'MY_API_KEY'" in str(warnings[0])
+
+
+def test_serialize_node__with_non_exist_code_input_path():
+    # GIVEN a code node with a non-existent code input path
+    class MyNode(CodeExecutionNode):
+        filepath = "non_existent_file.py"
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    with pytest.raises(NodeValidationError) as exc_info:
+        workflow_display.serialize()
+    assert "Filepath 'non_existent_file.py' does not exist" in str(exc_info.value)
+
+
+def test_serialize_node__with_environment_variable_references():
+    """
+    Tests that environment variable references in code node inputs serialize correctly.
+    """
+
+    # GIVEN a code node with environment variable references in code_inputs
+    class MyCodeExecutionNode(CodeExecutionNode):
+        code_inputs = {
+            "api_key": EnvironmentVariables.get("MY_API_KEY"),
+            "other_config": {"nested_key": EnvironmentVariables.get("NESTED_KEY")},
+        }
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyCodeExecutionNode
+
+    # 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 environment variable references
+    my_code_execution_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "CODE_EXECUTION"
+    )
+
+    # AND the api_key input should be serialized as an ENVIRONMENT_VARIABLE
+    api_key_input = next(inp for inp in my_code_execution_node["inputs"] if inp["key"] == "api_key")
+    assert api_key_input["value"] == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "type": "ENVIRONMENT_VARIABLE",
+                "data": {
+                    "environment_variable": "MY_API_KEY",
+                },
+            }
+        ],
+    }
+
+    # AND the nested environment variable should also be serialized correctly
+    other_config_input = next(inp for inp in my_code_execution_node["inputs"] if inp["key"] == "other_config")
+    assert other_config_input["value"]["combinator"] == "OR"
+    assert len(other_config_input["value"]["rules"]) == 1
+    assert other_config_input["value"]["rules"][0]["type"] == "CONSTANT_VALUE"
+    nested_data = other_config_input["value"]["rules"][0]["data"]
+    assert nested_data["type"] == "JSON"
+    assert nested_data["value"]["type"] == "DICTIONARY_REFERENCE"
+    assert len(nested_data["value"]["entries"]) == 1
+    assert nested_data["value"]["entries"][0]["key"] == "nested_key"
+    assert nested_data["value"]["entries"][0]["value"]["type"] == "ENVIRONMENT_VARIABLE"
+    assert nested_data["value"]["entries"][0]["value"]["environment_variable"] == "NESTED_KEY"
+
+
+def test_serialize_node__with_non_exist_code_input_path_with_dry_run():
+    # GIVEN a code node with a non-existent code input path
+    class MyNode(CodeExecutionNode):
+        filepath = "non_existent_file.py"
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow with dry_run=True
+    workflow_display = get_workflow_display(workflow_class=Workflow, dry_run=True)
+    data: dict = workflow_display.serialize()
+
+    # THEN the workflow should not raise an error
+    assert data == {
+        "workflow_raw_data": {
+            "nodes": [
+                {
+                    "id": "9b9e2a5d-01a4-46b2-80a3-d9484b2c0e08",
+                    "type": "ENTRYPOINT",
+                    "inputs": [],
+                    "data": {"label": "Entrypoint Node", "source_handle_id": "3e2a3f52-5047-4e2e-9a21-37bd43c63250"},
+                    "display_data": {"position": {"x": 0.0, "y": -50.0}},
+                    "base": None,
+                    "definition": None,
+                },
+                {
+                    "id": "f41cebba-a048-4852-a3d7-0f3100927166",
+                    "type": "CODE_EXECUTION",
+                    "inputs": [
+                        {
+                            "id": "b807bd5f-7d49-4597-89cf-044aa84cf8d6",
+                            "key": "code",
+                            "value": {
+                                "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "JSON", "value": None}}],
+                                "combinator": "OR",
+                            },
+                        },
+                        {
+                            "id": "dcd1e1a4-49e2-4d6b-95f9-8ec5fd2f8f5d",
+                            "key": "runtime",
+                            "value": {
+                                "rules": [
+                                    {"type": "CONSTANT_VALUE", "data": {"type": "STRING", "value": "PYTHON_3_11_6"}}
+                                ],
+                                "combinator": "OR",
+                            },
+                        },
+                    ],
+                    "data": {
+                        "label": "My Node",
+                        "error_output_id": None,
+                        "source_handle_id": "dc9edb2e-4392-4a2c-ab92-cc1b9c0cbd53",
+                        "target_handle_id": "66e7ef63-518b-40e7-911a-e38e8bcaec81",
+                        "code_input_id": "b807bd5f-7d49-4597-89cf-044aa84cf8d6",
+                        "runtime_input_id": "dcd1e1a4-49e2-4d6b-95f9-8ec5fd2f8f5d",
+                        "output_type": "STRING",
+                        "packages": [],
+                        "output_id": "98cae9b9-45cc-4897-a0f5-df250b56c00d",
+                        "log_output_id": "66c06c97-a9d1-4abf-840f-3f6c29709612",
+                    },
+                    "display_data": {"position": {"x": 200.0, "y": -50.0}},
+                    "base": {
+                        "name": "CodeExecutionNode",
+                        "module": ["vellum", "workflows", "nodes", "displayable", "code_execution_node", "node"],
+                    },
+                    "definition": {
+                        "name": "MyNode",
+                        "module": [
+                            "vellum_ee",
+                            "workflows",
+                            "display",
+                            "nodes",
+                            "vellum",
+                            "tests",
+                            "test_code_execution_node",
+                        ],
+                    },
+                    "ports": [{"id": "dc9edb2e-4392-4a2c-ab92-cc1b9c0cbd53", "name": "default", "type": "DEFAULT"}],
+                    "trigger": {"id": "66e7ef63-518b-40e7-911a-e38e8bcaec81", "merge_behavior": "AWAIT_ANY"},
+                    "outputs": [
+                        {
+                            "id": "98cae9b9-45cc-4897-a0f5-df250b56c00d",
+                            "name": "result",
+                            "schema": {"type": "string"},
+                            "type": "STRING",
+                            "value": None,
+                        },
+                        {
+                            "id": "66c06c97-a9d1-4abf-840f-3f6c29709612",
+                            "name": "log",
+                            "schema": {"type": "string"},
+                            "type": "STRING",
+                            "value": None,
+                        },
+                    ],
+                },
+            ],
+            "edges": [
+                {
+                    "id": "85e0961e-f968-49a0-beed-e21373f0ecda",
+                    "source_node_id": "9b9e2a5d-01a4-46b2-80a3-d9484b2c0e08",
+                    "source_handle_id": "3e2a3f52-5047-4e2e-9a21-37bd43c63250",
+                    "target_node_id": "f41cebba-a048-4852-a3d7-0f3100927166",
+                    "target_handle_id": "66e7ef63-518b-40e7-911a-e38e8bcaec81",
+                    "type": "DEFAULT",
+                }
+            ],
+            "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
+            "definition": {
+                "name": "Workflow",
+                "module": ["vellum_ee", "workflows", "display", "nodes", "vellum", "tests", "test_code_execution_node"],
+            },
+            "output_values": [],
+        },
+        "input_variables": [],
+        "state_variables": [],
+        "output_variables": [],
+    }
+
+
+def test_serialize_node__with_custom_output_type():
+    # GIVEN a code node with a custom output type
+    class MyNode(CodeExecutionNode[BaseState, dict[str, int]]):
+        code = """\
+return {
+    "hello": 1,
+}
+"""
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node's outputs should serialize correctly
+    my_code_execution_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "CODE_EXECUTION"
+    )
+    assert my_code_execution_node["outputs"] == [
+        {
+            "id": "01de8e8b-5e0e-4344-93b0-e002bbaed840",
+            "name": "result",
+            "value": None,
+            "type": "JSON",
+            "schema": {"type": "object", "additionalProperties": {"type": "integer"}},
+        },
+        {
+            "id": "64bf62e0-adc7-48cc-b689-8bf9b7e4eeef",
+            "name": "log",
+            "value": None,
+            "type": "STRING",
+            "schema": {"type": "string"},
+        },
+    ]
+    assert my_code_execution_node["data"]["output_type"] == "JSON"