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

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

@@ -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 → vellum_ai-1.1.5}/ee/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_ai-1.1.4 → vellum_ai-1.1.5}/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.1.5}/ee/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_ai-1.1.4 → vellum_ai-1.1.5}/ee/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_ai-1.1.4 → vellum_ai-1.1.5}/ee/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_ai-1.1.4 → vellum_ai-1.1.5}/ee/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_ai-1.1.4 → vellum_ai-1.1.5}/ee/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_ai-1.1.4 → vellum_ai-1.1.5}/ee/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

{vellum_ai-1.1.4 → vellum_ai-1.1.5}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "1.1.4"
+version = "1.1.5"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-1.1.4 → vellum_ai-1.1.5}/src/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_ai-1.1.4 → vellum_ai-1.1.5}/src/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_ai-1.1.4 → vellum_ai-1.1.5}/src/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_ai-1.1.4 → vellum_ai-1.1.5}/src/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