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

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -4,13 +4,12 @@ import sys
 import traceback
 from typing import Any, Optional, Tuple, Union
 
-from pydantic import BaseModel
-
+from vellum import Vellum
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
-from vellum.workflows.nodes.utils import cast_to_output_type
+from vellum.workflows.nodes.utils import cast_to_output_type, wrap_inputs_for_backward_compatibility
 from vellum.workflows.state.context import WorkflowContext
-from vellum.workflows.types.code_execution_node_wrappers import ListWrapper, clean_for_dict_wrapper
 from vellum.workflows.types.core import EntityInputsInterface
 
 
@@ -43,6 +42,7 @@ def run_code_inline(
     inputs: EntityInputsInterface,
     output_type: Any,
     filepath: str,
+    vellum_client: Vellum,
 ) -> Tuple[str, Any]:
     log_buffer = io.StringIO()
 
@@ -51,27 +51,14 @@ def run_code_inline(
         print_line = f"{' '.join(str_args)}\n"
         log_buffer.write(print_line)
 
-    def wrap_value(value):
-        if isinstance(value, list):
-            return ListWrapper(
-                [
-                    # Convert VellumValue to dict with its fields
-                    (
-                        item.model_dump()
-                        if isinstance(item, BaseModel)
-                        else clean_for_dict_wrapper(item) if isinstance(item, (dict, list, str)) else item
-                    )
-                    for item in value
-                ]
-            )
-        return clean_for_dict_wrapper(value)
-
+    wrapped_inputs = wrap_inputs_for_backward_compatibility(inputs)
     exec_globals = {
-        "__arg__inputs": {name: wrap_value(value) for name, value in inputs.items()},
+        "__arg__inputs": wrapped_inputs,
         "__arg__out": None,
         "print": _inline_print,
+        "vellum_client": vellum_client,
     }
-    run_args = [f"{name}=__arg__inputs['{name}']" for name in inputs.keys()]
+    run_args = [f"{name}=__arg__inputs['{name}']" for name, value in inputs.items() if value is not undefined]
     execution_code = f"""\
 {code}
 

vellum/workflows/nodes/displayable/conditional_node/node.py

@@ -1,24 +1,28 @@
-from typing import Set
+from typing import Generic, Set
 
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
-from vellum.workflows.ports.node_ports import NodePorts
 from vellum.workflows.ports.port import Port
 from vellum.workflows.ports.utils import validate_ports
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.generics import StateType
 
 
-class ConditionalNode(BaseNode):
+class ConditionalNode(BaseNode[StateType], Generic[StateType]):
     """
     Used to conditionally determine which port to invoke next. This node exists to be backwards compatible with
     Vellum's Conditional Node, and for most cases, you should extend `BaseNode.Ports` directly.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:split"
+        color = "corn"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
-    class Ports(NodePorts):
+    class Ports(BaseNode.Ports):
         def __call__(self, outputs: BaseOutputs, state: BaseState) -> Set[Port]:
             all_ports = [port for port in self.__class__]
             enforce_single_invoked_port = validate_ports(all_ports)

vellum/workflows/nodes/displayable/final_output_node/node.py

@@ -1,12 +1,15 @@
 from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args
 
 from vellum.workflows.constants import undefined
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
 from vellum.workflows.nodes.utils import cast_to_output_type
+from vellum.workflows.ports import NodePorts
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.validate import validate_target_type
 
 _OutputType = TypeVar("_OutputType")
 
@@ -26,6 +29,7 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
             **annotations,
             "value": parent.get_output_type(),
         }
+
         return parent
 
     def get_output_type(cls) -> Type:
@@ -44,9 +48,16 @@ class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], meta
     This provides backward compatibility with Vellum's Final Output Node.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:circle-stop"
+        color = "teal"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
+    class Ports(NodePorts):
+        pass
+
     class Outputs(BaseNode.Outputs):
         # We use our mypy plugin to override the _OutputType with the actual output type
         # for downstream references to this output.
@@ -61,3 +72,26 @@ class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], meta
                 self.__class__.get_output_type(),
             )
         )
+
+    __simulates_workflow_output__ = True
+
+    @classmethod
+    def __validate__(cls) -> None:
+        cls._validate_output_type_consistency()
+
+    @classmethod
+    def _validate_output_type_consistency(cls) -> None:
+        """
+        Validates that the declared output type of FinalOutputNode matches
+        the type of the descriptor assigned to the 'value' attribute in its Outputs class.
+
+        Raises ValueError if there's a type mismatch.
+        """
+        declared_output_type = cls.get_output_type()
+        value_descriptor = cls.Outputs.value.instance
+
+        if isinstance(value_descriptor, BaseDescriptor):
+            try:
+                validate_target_type(declared_output_type, value_descriptor.normalized_type)
+            except ValueError as e:
+                raise ValueError(f"Failed to validate output type for node '{cls.__name__}': {e}")

vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py

@@ -1,11 +1,18 @@
 import pytest
+from typing import Any, Dict
 
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.core.map_node import MapNode
+from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.nodes.displayable.inline_prompt_node import InlinePromptNode
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import Json
+from vellum.workflows.workflows.base import BaseWorkflow
 
 
-def test_final_output_node__mismatched_output_type():
+def test_final_output_node__mismatched_output_type_should_raise_exception_when_ran():
     # GIVEN a FinalOutputNode with a mismatched output type
     class StringOutputNode(FinalOutputNode[BaseState, str]):
         class Outputs(FinalOutputNode.Outputs):
@@ -18,3 +25,141 @@ def test_final_output_node__mismatched_output_type():
 
     # THEN an error is raised
     assert str(exc_info.value) == "Expected an output of type 'str', but received 'dict'"
+
+
+def test_final_output_node__mismatched_output_type_should_raise_exception():
+    # GIVEN a FinalOutputNode declared with list output type but has a string value type
+    class Output(FinalOutputNode[BaseState, list]):
+        """Output the extracted invoice line items as an array of objects."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = InlinePromptNode.Outputs.text
+
+    # WHEN attempting to validate the node class
+    # THEN a ValueError should be raised during validation
+    with pytest.raises(ValueError) as exc_info:
+        Output.__validate__()
+
+    # AND the error message should indicate the type mismatch
+    assert (
+        str(exc_info.value)
+        == "Failed to validate output type for node 'Output': Output type mismatch: declared type 'list' but "
+        "the 'value' Output has type(s) 'str'. "
+    )
+
+
+def test_final_output_node__mismatched_output_type_in_state_should_raise_exception():
+    # GIVEN a state with a str type
+    class State(BaseState):
+        foo: str
+
+    # AND a FinalOutputNode declared with list output type
+    class Output(FinalOutputNode[BaseState, list]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = State.foo
+
+    # WHEN attempting to validate the node class
+    # THEN a ValueError should be raised during validation
+    with pytest.raises(ValueError) as exc_info:
+        Output.__validate__()
+
+    # AND the error message should indicate the type mismatch
+    assert (
+        str(exc_info.value)
+        == "Failed to validate output type for node 'Output': Output type mismatch: declared type 'list' but "
+        "the 'value' Output has type(s) 'str'. "
+    )
+
+
+def test_final_output_node__matching_output_type_should_pass_validation():
+    # GIVEN a FinalOutputNode declared with correct matching types
+    class CorrectOutput(FinalOutputNode[BaseState, str]):
+        """Output with correct type matching."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = InlinePromptNode.Outputs.text
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        CorrectOutput.__validate__()
+    except ValueError:
+        pytest.fail("Validation should not raise an exception for correct type matching")
+
+
+def test_final_output_node__dict_and_Dict_should_be_compatible():
+    """
+    Tests that FinalOutputNode validation recognizes dict and Dict[str, Any] as compatible types.
+    """
+
+    # GIVEN a FinalOutputNode declared with dict output type
+    # AND the value descriptor has Dict[str, Any] type
+    class DictOutputNode(FinalOutputNode[BaseState, dict]):
+        """Output with dict type."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = OutputReference(
+                name="value",
+                types=(Dict[str, Any],),
+                instance=None,
+                outputs_class=FinalOutputNode.Outputs,
+            )
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        DictOutputNode.__validate__()
+    except ValueError as e:
+        pytest.fail(f"Validation should not raise an exception for dict/Dict compatibility: {e}")
+
+
+def test_final_output_node__any_output_type_should_accept_json():
+    """
+    Tests that FinalOutputNode with Any output type accepts a TemplatingNode with Json output type.
+    """
+
+    # GIVEN a TemplatingNode with Json output type
+    class JsonTemplatingNode(TemplatingNode[BaseState, Json]):
+        """Templating node that outputs Json."""
+
+        template = '{"key": "value"}'
+
+    # AND a FinalOutputNode with Any output type referencing the TemplatingNode
+    class AnyOutputNode(FinalOutputNode[BaseState, Any]):
+        """Output with Any type."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = JsonTemplatingNode.Outputs.result
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        AnyOutputNode.__validate__()
+    except ValueError as e:
+        pytest.fail(f"Validation should not raise an exception when Any accepts Json: {e}")
+
+
+def test_final_output_node__list_str_output_type_should_pass_validation():
+    """
+    Tests that FinalOutputNode with list[str] output type accepts a descriptor with List[str] type.
+    """
+
+    # GIVEN value descriptor has List[str] type
+    class MySubworkflow(BaseWorkflow):
+        class Outputs(BaseWorkflow.Outputs):
+            result: str
+
+    class MyMap(MapNode):
+        subworkflow = MySubworkflow
+
+    # AND a FinalOutputNode declared with list[str] output type
+    class ListStrOutputNode(FinalOutputNode[BaseState, list[str]]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = MyMap.Outputs.result
+
+    # WHEN attempting to validate the node class
+    try:
+        ListStrOutputNode.__validate__()
+    except Exception as e:
+        # THEN validation should pass without raising an exception
+        pytest.fail(f"Validation should not raise an exception for list[str]/List[str] compatibility: {e}")

vellum/workflows/nodes/displayable/guardrail_node/node.py

@@ -3,12 +3,11 @@ from typing import Any, ClassVar, Dict, Generic, List, Optional, Union, cast
 
 from vellum import ChatHistoryInput, ChatMessage, JsonInput, MetricDefinitionInput, NumberInput, StringInput
 from vellum.client import ApiError
-from vellum.core import RequestOptions
+from vellum.client.core import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.generics import StateType
 
@@ -23,6 +22,10 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
     request_options: Optional[RequestOptions] - The request options to use for the Metric
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:shield-check"
+        color = "rose"
+
     metric_definition: ClassVar[Union[UUID, str]]
 
     metric_inputs: ClassVar[EntityInputsInterface] = {}
@@ -33,7 +36,7 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
-    class Outputs(BaseOutputs):
+    class Outputs(BaseNode.Outputs):
         score: float
         normalized_score: Optional[float]
         log: Optional[str]
@@ -98,7 +101,13 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
         else:
             reason = None
 
-        return self.Outputs(score=score, normalized_score=normalized_score, log=log, reason=reason, **metric_outputs)
+        return self.Outputs(
+            score=score,
+            normalized_score=normalized_score,
+            log=log,
+            reason=reason,
+            **metric_outputs,  # type: ignore [arg-type]
+        )
 
     def _compile_metric_inputs(self) -> List[MetricDefinitionInput]:
         # TODO: We may want to consolidate with prompt deployment input compilation

vellum/workflows/nodes/displayable/inline_prompt_node/node.py

@@ -1,16 +1,16 @@
-import json
-from typing import Any, Dict, Iterator, Type, Union
+from typing import Any, Dict, Generic, Iterator, Type, Union
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BaseInlinePromptNode as BaseInlinePromptNode
+from vellum.workflows.nodes.displayable.bases.utils import process_additional_prompt_outputs
 from vellum.workflows.outputs import BaseOutput
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 
 
-class InlinePromptNode(BaseInlinePromptNode[StateType]):
+class InlinePromptNode(BaseInlinePromptNode[StateType], Generic[StateType]):
     """
     Used to execute a Prompt defined inline.
 
@@ -23,6 +23,10 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
     request_options: Optional[RequestOptions] - The request options to use for the Prompt Execution
     """
 
+    class Display(BaseInlinePromptNode.Display):
+        icon = "vellum:icon:text-size"
+        color = "navy"
+
     class Trigger(BaseInlinePromptNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
@@ -45,39 +49,8 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
                 code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
-        string_outputs = []
-        json_output = None
-
-        should_parse_json = False
-        if hasattr(self, "parameters"):
-            custom_params = self.parameters.custom_parameters
-            if custom_params and isinstance(custom_params, dict):
-                json_schema = custom_params.get("json_schema", {})
-                if (isinstance(json_schema, dict) and "schema" in json_schema) or custom_params.get("json_mode", {}):
-                    should_parse_json = True
-
-        for output in outputs:
-            if output.value is None:
-                continue
-
-            if output.type == "STRING":
-                string_outputs.append(output.value)
-                if should_parse_json:
-                    try:
-                        parsed_json = json.loads(output.value)
-                        json_output = parsed_json
-                    except (json.JSONDecodeError, TypeError):
-                        pass
-            elif output.type == "JSON":
-                string_outputs.append(json.dumps(output.value, indent=4))
-                json_output = output.value
-            elif output.type == "FUNCTION_CALL":
-                string_outputs.append(output.value.model_dump_json(indent=4))
-            else:
-                string_outputs.append(output.value.message)
-
-        value = "\n".join(string_outputs)
-        yield BaseOutput(name="text", value=value)
+        text_output, json_output = process_additional_prompt_outputs(outputs)
+        yield BaseOutput(name="text", value=text_output)
 
         if json_output:
             yield BaseOutput(name="json", value=json_output)

vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py

@@ -197,6 +197,9 @@ def test_inline_prompt_node__api_error__invalid_inputs_node_exception(
     assert e.value.code == expected_code
     assert e.value.message == expected_message
 
+    # AND the node exception includes the response body in raw_data
+    assert e.value.raw_data == exception.body
+
 
 def test_inline_prompt_node__chat_history_inputs(vellum_adhoc_prompt_client):
     # GIVEN a prompt node with a chat history input

vellum/workflows/nodes/displayable/merge_node/node.py

@@ -1,12 +1,19 @@
+from typing import Generic
+
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.generics import StateType
 
 
-class MergeNode(BaseNode):
+class MergeNode(BaseNode[StateType], Generic[StateType]):
     """
     Used to merge the control flow of multiple nodes into a single node. This node exists to be backwards compatible
     with Vellum's Merge Node, and for most cases, you should extend from `BaseNode.Trigger` directly.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:merge"
+        color = "tomato"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY

vellum/workflows/nodes/displayable/note_node/node.py

@@ -1,12 +1,22 @@
+from typing import Any, Dict, Generic, Optional
+
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.generics import StateType
 
 
-class NoteNode(BaseNode):
+class NoteNode(BaseNode[StateType], Generic[StateType]):
     """
     A no-op Node purely used to display a note in the Vellum UI.
     """
 
+    text: str = ""
+    style: Optional[Dict[str, Any]] = None
+
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:note"
+        color = "cyan"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 

vellum/workflows/nodes/displayable/prompt_deployment_node/node.py

@@ -1,10 +1,10 @@
-import json
 from typing import Any, Dict, Iterator, Type, Union
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BasePromptDeploymentNode as BasePromptDeploymentNode
+from vellum.workflows.nodes.displayable.bases.utils import process_additional_prompt_outputs
 from vellum.workflows.outputs import BaseOutput
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
@@ -26,6 +26,10 @@ class PromptDeploymentNode(BasePromptDeploymentNode[StateType]):
     request_options: Optional[RequestOptions] - The request options to use for the Prompt Execution
     """
 
+    class Display(BasePromptDeploymentNode.Display):
+        icon = "vellum:icon:text-size"
+        color = "navy"
+
     class Trigger(BasePromptDeploymentNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
@@ -48,28 +52,8 @@ class PromptDeploymentNode(BasePromptDeploymentNode[StateType]):
                 code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
-        string_outputs = []
-        json_output = None
-
-        for output in outputs:
-            if output.value is None:
-                continue
-
-            if output.type == "STRING":
-                string_outputs.append(output.value)
-                try:
-                    json_output = json.loads(output.value)
-                except (json.JSONDecodeError, TypeError):
-                    pass
-            elif output.type == "JSON":
-                string_outputs.append(json.dumps(output.value, indent=4))
-            elif output.type == "FUNCTION_CALL":
-                string_outputs.append(output.value.model_dump_json(indent=4))
-            else:
-                string_outputs.append(output.value.message)
-
-        value = "\n".join(string_outputs)
-        yield BaseOutput(name="text", value=value)
+        text_output, json_output = process_additional_prompt_outputs(outputs)
+        yield BaseOutput(name="text", value=text_output)
 
         if json_output:
             yield BaseOutput(name="json", value=json_output)

vellum/workflows/nodes/displayable/search_node/node.py

@@ -1,8 +1,11 @@
 import json
 from typing import ClassVar
 
+from vellum.utils.json_encoder import VellumJsonEncoder
+from vellum.workflows.constants import undefined
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BaseSearchNode as BaseSearchNode
-from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 
@@ -19,6 +22,10 @@ class SearchNode(BaseSearchNode[StateType]):
     chunk_separator: str = "\n\n#####\n\n" - The separator to use when joining the text of each search result
     """
 
+    class Display(BaseSearchNode.Display):
+        icon = "vellum:icon:magnifying-glass-waveform"
+        color = "purple"
+
     chunk_separator: ClassVar[str] = "\n\n#####\n\n"
 
     class Trigger(BaseSearchNode.Trigger):
@@ -35,8 +42,14 @@ class SearchNode(BaseSearchNode[StateType]):
         text: str
 
     def run(self) -> Outputs:
+        if self.query is undefined or self.query is None or self.query == "":
+            raise NodeException(
+                message="Search query is required but was not provided",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+
         if not isinstance(self.query, str):
-            self.query = json.dumps(self.query, cls=DefaultStateEncoder)
+            self.query = json.dumps(self.query, cls=VellumJsonEncoder)
 
         results = self._perform_search().results
         text = self.chunk_separator.join([r.text for r in results])

vellum/workflows/nodes/displayable/search_node/tests/test_node.py

@@ -1,3 +1,4 @@
+import pytest
 import json
 
 from vellum import SearchResponse, SearchResult, SearchResultDocument
@@ -10,6 +11,8 @@ from vellum.client.types.search_weights_request import SearchWeightsRequest
 from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
 from vellum.client.types.vellum_value_logical_condition_group_request import VellumValueLogicalConditionGroupRequest
 from vellum.client.types.vellum_value_logical_condition_request import VellumValueLogicalConditionRequest
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.types import (
     MetadataLogicalCondition,
     MetadataLogicalConditionGroup,
@@ -215,3 +218,33 @@ def test_run_workflow__chat_history_as_query(vellum_client):
     assert json.loads(vellum_client.search.call_args.kwargs["query"]) == [
         {"role": "USER", "text": "Hello, world!", "source": None, "content": None}
     ]
+
+
+@pytest.mark.parametrize("invalid_query", [None, ""])
+def test_run_workflow__invalid_query_raises_validation_error(invalid_query):
+    """Confirm that missing/None/empty query raises proper user-facing validation error"""
+
+    class MySearchNode(SearchNode):
+        query = invalid_query  # type: ignore[assignment]
+        document_index = "document_index"
+
+    with pytest.raises(NodeException) as exc_info:
+        MySearchNode().run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "query" in exc_info.value.message.lower()
+    assert "required" in exc_info.value.message.lower() or "missing" in exc_info.value.message.lower()
+
+
+def test_run_workflow__missing_query_attribute_raises_validation_error():
+    """Confirm that a SearchNode without a query attribute defined raises INVALID_INPUTS"""
+
+    class MySearchNode(SearchNode):
+        document_index = "document_index"
+
+    with pytest.raises(NodeException) as exc_info:
+        MySearchNode().run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "query" in exc_info.value.message.lower()
+    assert "required" in exc_info.value.message.lower()

vellum/workflows/nodes/displayable/set_state_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import SetStateNode
+
+__all__ = [
+    "SetStateNode",
+]