vellum-ai 1.1.4__py3-none-any.whl → 1.1.5__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -27,10 +27,10 @@ class BaseClientWrapper:
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
-            "User-Agent": "vellum-ai/1.1.4",
+            "User-Agent": "vellum-ai/1.1.5",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "1.1.4",
+            "X-Fern-SDK-Version": "1.1.5",
             **(self.get_custom_headers() or {}),
         }
         if self._api_version is not None:

vellum/workflows/emitters/vellum_emitter.py

@@ -11,6 +11,8 @@ from vellum.workflows.state.base import BaseState
 
 logger = logging.getLogger(__name__)
 
+DISALLOWED_EVENTS = {"workflow.execution.streaming", "node.execution.streaming"}
+
 
 class VellumEmitter(BaseWorkflowEmitter):
     """
@@ -53,6 +55,9 @@ class VellumEmitter(BaseWorkflowEmitter):
         if not self._context:
             return
 
+        if event.name in DISALLOWED_EVENTS:
+            return
+
         try:
             event_data = default_serializer(event)
 

vellum/workflows/state/context.py

@@ -1,12 +1,10 @@
 from functools import cached_property
-import inspect
 from queue import Queue
 from uuid import uuid4
 from typing import TYPE_CHECKING, Dict, List, Optional, Type
 
 from vellum import Vellum
-from vellum.workflows.context import ExecutionContext, get_execution_context, set_execution_context
-from vellum.workflows.emitters.vellum_emitter import VellumEmitter
+from vellum.workflows.context import ExecutionContext, get_execution_context
 from vellum.workflows.events.types import ExternalParentContext
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
@@ -14,7 +12,6 @@ from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.vellum_client import create_vellum_client
 
 if TYPE_CHECKING:
-    from vellum.workflows.emitters.base import BaseWorkflowEmitter
     from vellum.workflows.events.workflow import WorkflowEvent
 
 
@@ -29,36 +26,15 @@ class WorkflowContext:
         self._vellum_client = vellum_client
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
         self._node_output_mocks_map: Dict[Type[BaseOutputs], List[MockNodeExecution]] = {}
-        # Clone the current thread-local execution context to avoid mutating global state
-        current_execution_context = get_execution_context()
-
-        # Resolve parent_context preference: provided > current > new external
-        resolved_parent_context = (
-            execution_context.parent_context
-            if execution_context is not None and execution_context.parent_context is not None
-            else current_execution_context.parent_context
-        )
-        if resolved_parent_context is None:
-            resolved_parent_context = ExternalParentContext(span_id=uuid4())
-
-        # Resolve trace_id preference: provided (if set) > current (if set) > new uuid
-        if execution_context is not None and int(execution_context.trace_id) != 0:
-            resolved_trace_id = execution_context.trace_id
-        elif int(current_execution_context.trace_id) != 0:
-            resolved_trace_id = current_execution_context.trace_id
-        else:
-            resolved_trace_id = uuid4()
+        self._execution_context = get_execution_context()
 
-        # Construct a single, resolved execution context for this workflow instance
-        self._execution_context = ExecutionContext(
-            parent_context=resolved_parent_context,
-            trace_id=resolved_trace_id,
-        )
+        if execution_context is not None:
+            self._execution_context.trace_id = execution_context.trace_id
+            if execution_context.parent_context is not None:
+                self._execution_context.parent_context = execution_context.parent_context
 
-        # Ensure the thread-local context has a parent_context for nodes that read it directly
-        if current_execution_context.parent_context is None:
-            current_execution_context.parent_context = resolved_parent_context
-            set_execution_context(current_execution_context)
+        if self._execution_context.parent_context is None:
+            self._execution_context.parent_context = ExternalParentContext(span_id=uuid4())
 
         self._generated_files = generated_files
 
@@ -128,26 +104,6 @@ class WorkflowContext:
             # For custom domains, assume the same pattern: api.* -> app.*
             return api_url.replace("api.", "app.", 1)
 
-    def get_emitters_for_workflow(self) -> List["BaseWorkflowEmitter"]:
-        """
-        Get the default emitters that should be attached to workflows using this context.
-
-        Returns:
-            List of emitters, including VellumEmitter if monitoring is enabled.
-        """
-        try:
-            frame = inspect.currentframe()
-            caller = frame.f_back if frame else None
-            if caller and "self" in caller.f_locals:
-                workflow_instance = caller.f_locals["self"]
-                class_level_emitters = getattr(workflow_instance.__class__, "emitters", None)
-                if isinstance(class_level_emitters, list) and len(class_level_emitters) > 0:
-                    return class_level_emitters
-        except Exception:
-            pass
-
-        return [VellumEmitter()]
-
     def _emit_subworkflow_event(self, event: "WorkflowEvent") -> None:
         if self._event_queue:
             self._event_queue.put(event)

vellum/workflows/workflows/base.py

@@ -241,7 +241,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
     ):
         self._parent_state = parent_state
         self._context = context or WorkflowContext()
-        self.emitters = emitters or self._context.get_emitters_for_workflow()
+        self.emitters = emitters or (self.emitters if hasattr(self, "emitters") else [])
         self.resolvers = resolvers or (self.resolvers if hasattr(self, "resolvers") else [])
         self._store = store or Store()
         self._execution_context = self._context.execution_context

{vellum_ai-1.1.4.dist-info → vellum_ai-1.1.5.dist-info}/METADATA

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

{vellum_ai-1.1.4.dist-info → vellum_ai-1.1.5.dist-info}/RECORD

@@ -28,10 +28,10 @@ vellum_ee/workflows/display/editor/__init__.py,sha256=MSAgY91xCEg2neH5d8jXx5wRdR
 vellum_ee/workflows/display/editor/types.py,sha256=x-tOOCJ6CF4HmiKDfCmcc3bOVfc1EBlP5o6u5WEfLoY,567
 vellum_ee/workflows/display/exceptions.py,sha256=Oys39dHoW-s-1dnlRSZxTntMq8_macj-b2CT_6dqzJs,355
 vellum_ee/workflows/display/nodes/__init__.py,sha256=jI1aPBQf8DkmrYoZ4O-wR1duqZByOf5mDFmo_wFJPE4,307
-vellum_ee/workflows/display/nodes/base_node_display.py,sha256=oJJiNmF9na2rS5PMMWZLwbajJkssEeBCESDmmz9jy2s,17933
+vellum_ee/workflows/display/nodes/base_node_display.py,sha256=EoAeFrfo8bk4ad0G6ZXeco8iRmfi6rLgvAEQ40Ald4U,17975
 vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=jI_kUi9LnNLDpY63QtlC4TfN8P571VN4LpzH0I1ZtLk,1149
 vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/nodes/tests/test_base_node_display.py,sha256=Z4Mf7xLCNiblSbpKI0BrV5modQr-ZcFzhfir_OSyTTs,2997
+vellum_ee/workflows/display/nodes/tests/test_base_node_display.py,sha256=wBoCqULS4XO3s9Vwhd9v4g10opfBFqeZgRqB8CoFz0c,3015
 vellum_ee/workflows/display/nodes/types.py,sha256=St1BB6no528OyELGiyRabWao0GGw6mLhstQAvEACbGk,247
 vellum_ee/workflows/display/nodes/utils.py,sha256=sloya5TpXsnot1HURc9L51INwflRqUzHxRVnCS9Cd-4,973
 vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=nUIgH2s0-7IbQRNrBhLPyRNe8YIrx3Yo9HeeW-aXXFk,1668
@@ -75,7 +75,7 @@ vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DE
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=Y-ajeT65b5varmrZCw6L3hir4hJCFq-eO0jZfRcrs7g,1886
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=aEmF6zpmmHdFgiG4V81EweOk5qO6cfukR782dYSE960,14060
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=DFNp_iUxlyE-zCJBH9ab3e1jQEK-NSE9M2CQd4NCjY8,24853
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=tZ_YK1NpMd8_2ADcbF3PKM2i7fIbSyjUMsP8MCWx2g0,26643
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=s6_mnk0pkztU59wYpSfOFpMhAJaRjmyfxM6WJGtnD4Y,6456
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=PkSgghJDz0fpDB72HHPjLjo8LkZk-HpUkCQzRLX-iVw,40611
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py,sha256=dsJr8I9AdPwMOGszirfNDzZP2Ychd94aAKuPXAzknMk,4632
@@ -96,14 +96,14 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py,sha256=V8b6gKghLlO7PJI8xeNdnfn8aII0W_IFQvSQBQM62UQ,7721
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py,sha256=hDWtKXmGI1CKhTwTNqpu_d5RkE5n7SolMLtgd87KqTI,3856
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_composio_serialization.py,sha256=AUfULtIangNYkvLH3jd2Ar8X5ulW4tGmezeCfMmXFUU,3697
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_inline_workflow_serialization.py,sha256=4t1lkN2nsZF6lFqP6QnskUQWJlhasF8C2_f6atzk8ZY,26298
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_inline_workflow_serialization.py,sha256=K_Wx8QLFeigziK8KvbRFm5M_bBSmVzcsmQnagC97DWo,26363
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py,sha256=9ORex-okcpwbbkxEDJyyRlbPid6zLSDduK0fBfrp8kk,2415
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py,sha256=B0rDsCvO24qPp0gkmj8SdTDY5CxZYkvKwknsKBuAPyA,10017
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py,sha256=txiUpg0BrGJiP2hrTVpwoT4bGqoo0rImSFFvOpigX8c,10082
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_workflow_deployment_serialization.py,sha256=XIZZr5POo2NLn2uEWm9EC3rejeBMoO4X-JtzTH6mvp4,4074
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=pLCyMScV88DTBXRH7jXaXOEA1GBq8NIipCUFwIAWnwI,2771
 vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=J4ouI8KxbMfxQP2Zq_9cWMGYgbjCWmKzjCJEtnSJb0I,5829
 vellum_ee/workflows/display/tests/workflow_serialization/test_workflow_input_parameterization_error.py,sha256=vAdmn3YTBDpo55znbydQxsgg9ASqHcvsUPwiBR_7wfo,1461
-vellum_ee/workflows/display/types.py,sha256=qRSh7HstBjcAft6dIrij9dLZTanlRWhNwDFfggqJXzs,3174
+vellum_ee/workflows/display/types.py,sha256=cyZruu4sXAdHjwuFc7dydM4DcFNf-pp_CmulXItxac4,3679
 vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/utils/auto_layout.py,sha256=f4GiLn_LazweupfqTpubcdtdfE_vrOcmZudSsnYIY9E,3906
 vellum_ee/workflows/display/utils/exceptions.py,sha256=LSwwxCYNxFkf5XMUcFkaZKpQ13OSrI7y_bpEUwbKVk0,169
@@ -114,9 +114,9 @@ vellum_ee/workflows/display/utils/tests/test_auto_layout.py,sha256=vfXI769418s9v
 vellum_ee/workflows/display/utils/vellum.py,sha256=mtoXmSYwR7rvrq-d6CzCW_auaJXTct0Mi1F0xpRCiNQ,5627
 vellum_ee/workflows/display/vellum.py,sha256=J2mdJZ1sdLW535DDUkq_Vm8Z572vhuxHxVZF9deKSdk,391
 vellum_ee/workflows/display/workflows/__init__.py,sha256=JTB9ObEV3l4gGGdtfBHwVJtTTKC22uj-a-XjTVwXCyA,148
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=ysOq1_con5SvdaeCMeECWEVC1qJzSBTIKA_lJNU0zqQ,41412
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=mEJPUZIBPjPdYDUW1ZeWiMTLaDKfnNQQ7kvcwiTJ_ac,42041
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=gxz76AeCqgAZ9D2lZeTiZzxY9eMgn3qOSfVgiqYcOh8,2028
-vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=9CkgKbM_zgq0l2bwst1hLFGyOigpjsBO4ZU6qSMXND4,36027
+vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=pb7BTH-ivRnya1LQU3j-MApWk_m8POpPNOdD0oEK82A,37847
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=7JphJcSO3H85qiC2DpFfBWjC3JjrbRmoynBC6KKHVsA,2710
 vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -147,7 +147,7 @@ vellum/client/README.md,sha256=gxc7JlJRBrBZpN5LHa2ORxYTRHFLPnWmnIugN8pmQh4,5600
 vellum/client/__init__.py,sha256=NEcLUhKRIK0OGGXTSQbarRaacFkJWFmM5iY4F9lgIDI,72131
 vellum/client/core/__init__.py,sha256=lTcqUPXcx4112yLDd70RAPeqq6tu3eFMe1pKOqkW9JQ,1562
 vellum/client/core/api_error.py,sha256=44vPoTyWN59gonCIZMdzw7M1uspygiLnr3GNFOoVL2Q,614
-vellum/client/core/client_wrapper.py,sha256=6vp_q4aKby-T0ukw9WHtdpLBDqY6TQ7mpOcJmVsaiE4,2840
+vellum/client/core/client_wrapper.py,sha256=DpUpUgjS2uzkAZFSPBjxVkQxUTezxAyRpEypFJZvapU,2840
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=d4NNbX8XvXP32z8KpK2Xovv33nFfruIrpz0QWxlgpZk,2663
 vellum/client/core/force_multipart.py,sha256=awxh5MtcRYe74ehY8U76jzv6fYM_w_D3Rur7KQQzSDk,429
@@ -1598,7 +1598,7 @@ vellum/workflows/edges/__init__.py,sha256=wSkmAnz9xyi4vZwtDbKxwlplt2skD7n3NsxkvR
 vellum/workflows/edges/edge.py,sha256=N0SnY3gKVuxImPAdCbPMPlHJIXbkQ3fwq_LbJRvVMFc,677
 vellum/workflows/emitters/__init__.py,sha256=d9QFOI3eVg6rzpSFLvrjkDYXWikf1tcp3ruTRa2Boyc,143
 vellum/workflows/emitters/base.py,sha256=Tcp13VMB-GMwEJdl-6XTPckspdOdwpMgBx22-PcQxds,892
-vellum/workflows/emitters/vellum_emitter.py,sha256=xiVlqe3B5fTR-e48nuneeodN1tkrBt0QQa0o9JHdyWg,5077
+vellum/workflows/emitters/vellum_emitter.py,sha256=c84BNgXkpOZDZglRz26IhmQtLFwk3DnRHwqOmf4GHIY,5223
 vellum/workflows/environment/__init__.py,sha256=TJz0m9dwIs6YOwCTeuN0HHsU-ecyjc1OJXx4AFy83EQ,121
 vellum/workflows/environment/environment.py,sha256=Ck3RPKXJvtMGx_toqYQQQF-ZwXm5ijVwJpEPTeIJ4_Q,471
 vellum/workflows/errors/__init__.py,sha256=tWGPu5xyAU8gRb8_bl0fL7OfU3wxQ9UH6qVwy4X4P_Q,113
@@ -1805,7 +1805,7 @@ vellum/workflows/runner/runner.py,sha256=sADrp593ia4auJ8fIP6jc5na3b5WaTvvBdjpkTt
 vellum/workflows/sandbox.py,sha256=GVJzVjMuYzOBnSrboB0_6MMRZWBluAyQ2o7syeaeBd0,2235
 vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
 vellum/workflows/state/base.py,sha256=m9fCqbZn21GshCVCjJTD1dPZEQjFrsMXqlg7tM9fIwM,24283
-vellum/workflows/state/context.py,sha256=HQbapsGX42rCEn1GuD-jWMpvjrQhosPVJSnashKN2vc,7234
+vellum/workflows/state/context.py,sha256=i9tQKTB2b2SBpNObtb8skTv63lQcwUDW1ZQ0qDTnGJA,5181
 vellum/workflows/state/delta.py,sha256=7h8wR10lRCm15SykaPj-gSEvvsMjCwYLPsOx3nsvBQg,440
 vellum/workflows/state/encoder.py,sha256=9wj9za3pjLGaKh85IB71YvPDRDtHcw78wR5vyPyKWM8,2708
 vellum/workflows/state/store.py,sha256=uVe-oN73KwGV6M6YLhwZMMUQhzTQomsVfVnb8V91gVo,1147
@@ -1837,13 +1837,13 @@ vellum/workflows/utils/uuids.py,sha256=DFzPv9RCvsKhvdTEIQyfSek2A31D6S_QcmeLPbgrg
 vellum/workflows/utils/vellum_variables.py,sha256=vg7PLTqBoPLWXrcFR2vOfJpKkzvDC85cL2sSUqDWXWU,5503
 vellum/workflows/vellum_client.py,sha256=xkfoucodxNK5JR2-lbRqZx3xzDgExWkP6kySrpi_Ubc,1079
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
-vellum/workflows/workflows/base.py,sha256=FTW2y5O8bUgMD1aluavwQ-UYTURRW1DRxDwdynR0OZc,27343
+vellum/workflows/workflows/base.py,sha256=C17t-YpaBishP5PNThPgmMtYQsint55AcPIweBTVi7I,27354
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=ptMntHzVyy8ZuzNgeTuk7hREgKQ5UBdgq8VJFSGaW4Y,20832
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-1.1.4.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-1.1.4.dist-info/METADATA,sha256=vtS36kiUPHeBJrpddCNWGFQRpj7hTebtWkd4xCZVdKA,5547
-vellum_ai-1.1.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-1.1.4.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-1.1.4.dist-info/RECORD,,
+vellum_ai-1.1.5.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-1.1.5.dist-info/METADATA,sha256=tyN3xDCmqMqgJmFZweIhi6oIZXIrDFs3Xnac__QYrdI,5547
+vellum_ai-1.1.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-1.1.5.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-1.1.5.dist-info/RECORD,,

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -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_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_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -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_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_inline_workflow_serialization.py

@@ -44,7 +44,10 @@ def test_serialize_workflow():
         "type": "GENERIC",
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},
-            "comment": {"value": "\n    A tool calling node that calls the get_current_weather function.\n    "},
+            "comment": {
+                "expanded": True,
+                "value": "\n    A tool calling node that calls the get_current_weather function.\n    ",
+            },
         },
         "base": {
             "name": "ToolCallingNode",

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

@@ -44,7 +44,10 @@ def test_serialize_workflow():
         "type": "GENERIC",
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},
-            "comment": {"value": "\n    A tool calling node that calls the get_current_weather function.\n    "},
+            "comment": {
+                "expanded": True,
+                "value": "\n    A tool calling node that calls the get_current_weather function.\n    ",
+            },
         },
         "base": {
             "name": "ToolCallingNode",

vellum_ee/workflows/display/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, Iterator, List, Tuple, Type
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Type
 
 from vellum.client import Vellum as VellumClient
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -52,15 +52,26 @@ class WorkflowDisplayContext:
     edge_displays: EdgeDisplays = field(default_factory=dict)
     port_displays: PortDisplays = field(default_factory=dict)
     _errors: List[Exception] = field(default_factory=list)
+    _invalid_nodes: List[Type[BaseNode]] = field(default_factory=list)
     _dry_run: bool = False
 
-    def add_error(self, error: Exception) -> None:
+    def add_error(self, error: Exception, node: Optional[Type[BaseNode]] = None) -> None:
         if self._dry_run:
             self._errors.append(error)
             return
 
         raise error
 
+    def add_invalid_node(self, node: Type[BaseNode]) -> None:
+        """Track a node that failed to serialize."""
+        if node not in self._invalid_nodes:
+            self._invalid_nodes.append(node)
+
     @property
     def errors(self) -> Iterator[Exception]:
         return iter(self._errors)
+
+    @property
+    def invalid_nodes(self) -> Iterator[Type[BaseNode]]:
+        """Get an iterator over nodes that failed to serialize."""
+        return iter(self._invalid_nodes)

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -194,6 +194,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 serialized_node = node_display.serialize(self.display_context)
             except (NotImplementedError, NodeValidationError) as e:
                 self.display_context.add_error(e)
+                self.display_context.add_invalid_node(node)
                 continue
 
             serialized_nodes[node_display.node_id] = serialized_node
@@ -319,6 +320,10 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         # Add an edge for each edge in the workflow
         for target_node, entrypoint_display in self.display_context.entrypoint_displays.items():
             unadorned_target_node = get_unadorned_node(target_node)
+            # Skip edges to invalid nodes
+            if self._is_node_invalid(unadorned_target_node):
+                continue
+
             target_node_display = self.display_context.node_displays[unadorned_target_node]
             edges.append(
                 {
@@ -335,6 +340,12 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             unadorned_source_node_port = get_unadorned_port(source_node_port)
             unadorned_target_node = get_unadorned_node(target_node)
 
+            # Skip edges that reference invalid nodes
+            if self._is_node_invalid(unadorned_target_node) or self._is_node_invalid(
+                unadorned_source_node_port.node_class
+            ):
+                continue
+
             source_node_port_display = self.display_context.port_displays[unadorned_source_node_port]
             target_node_display = self.display_context.node_displays[unadorned_target_node]
 
@@ -939,5 +950,9 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         is_required = not has_default and not is_optional
         return is_required
 
+    def _is_node_invalid(self, node: Type[BaseNode]) -> bool:
+        """Check if a node failed to serialize and should be considered invalid."""
+        return node in self.display_context.invalid_nodes
+
 
 register_workflow_display_class(workflow_class=BaseWorkflow, workflow_display_class=BaseWorkflowDisplay)

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -1,6 +1,6 @@
 import pytest
-from uuid import uuid4
-from typing import Optional
+from uuid import UUID, uuid4
+from typing import Any, Optional
 
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
@@ -11,11 +11,13 @@ from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.displayable.final_output_node.node import FinalOutputNode
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import JsonObject
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.retry_node import BaseRetryNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.try_node import BaseTryNodeDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -955,3 +957,43 @@ def test_serialize_workflow__state_variables():
         "required": False,
         "extensions": {"color": None},
     }
+
+
+def test_serialize_workflow__with_complete_node_failure_prunes_edges():
+    """Test that edges are pruned when a node completely fails to serialize (serialized_node is null)."""
+
+    # GIVEN a node that completely fails to serialize
+    class FailingNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class FailingNodeDisplay(BaseNodeDisplay[FailingNode]):
+        def serialize(
+            self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        ) -> JsonObject:
+            raise NotImplementedError("Complete node serialization failure")
+
+    # AND a workflow with the failing node connected to another node
+    class WorkingNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            value: str
+
+    class Workflow(BaseWorkflow):
+        graph = FailingNode >> WorkingNode
+
+    # WHEN we serialize the workflow with dry_run=True
+    workflow_display = get_workflow_display(workflow_class=Workflow, dry_run=True)
+
+    # AND we register the failing display class
+    workflow_display.display_context.node_displays[FailingNode] = FailingNodeDisplay()
+
+    data: dict = workflow_display.serialize()
+
+    # THEN the workflow should serialize but with no edges (pruned due to invalid node)
+    assert data["workflow_raw_data"]["edges"] == []
+
+    # AND only the working node and entrypoint should be in the serialized nodes
+    assert len(data["workflow_raw_data"]["nodes"]) == 2
+    node_types = [node["type"] for node in data["workflow_raw_data"]["nodes"]]
+    assert "ENTRYPOINT" in node_types
+    assert "GENERIC" in node_types  # This is the WorkingNode that should still be serialized