vellum-ai 0.14.0__py3-none-any.whl → 0.14.2__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.0",
+            "X-Fern-SDK-Version": "0.14.2",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/prompts/blocks/compilation.py

@@ -1,5 +1,5 @@
 import json
-from typing import Optional, cast
+from typing import Sequence, Union, cast
 
 from vellum import (
     ChatMessage,
@@ -14,6 +14,7 @@ from vellum.client.types.audio_vellum_value import AudioVellumValue
 from vellum.client.types.function_call import FunctionCall
 from vellum.client.types.function_call_vellum_value import FunctionCallVellumValue
 from vellum.client.types.image_vellum_value import ImageVellumValue
+from vellum.client.types.number_input import NumberInput
 from vellum.client.types.vellum_audio import VellumAudio
 from vellum.client.types.vellum_image import VellumImage
 from vellum.prompts.blocks.exceptions import PromptCompilationError
@@ -22,15 +23,20 @@ from vellum.utils.templating.constants import DEFAULT_JINJA_CUSTOM_FILTERS
 from vellum.utils.templating.render import render_sandboxed_jinja_template
 from vellum.utils.typing import cast_not_optional
 
+PromptInput = Union[PromptRequestInput, NumberInput]
+
 
 def compile_prompt_blocks(
     blocks: list[PromptBlock],
-    inputs: list[PromptRequestInput],
+    inputs: Sequence[PromptInput],
     input_variables: list[VellumVariable],
 ) -> list[CompiledPromptBlock]:
     """Compiles a list of Prompt Blocks, performing all variable substitutions and Jinja templating needed."""
 
     sanitized_inputs = _sanitize_inputs(inputs)
+    inputs_by_name = {
+        (input_.name if isinstance(input_, NumberInput) else input_.key): input_ for input_ in sanitized_inputs
+    }
 
     compiled_blocks: list[CompiledPromptBlock] = []
     for block in blocks:
@@ -66,7 +72,7 @@ def compile_prompt_blocks(
 
             rendered_template = render_sandboxed_jinja_template(
                 template=block.template,
-                input_values={input_.key: input_.value for input_ in sanitized_inputs},
+                input_values={name: inp.value for name, inp in inputs_by_name.items()},
                 jinja_custom_filters=DEFAULT_JINJA_CUSTOM_FILTERS,
                 jinja_globals=DEFAULT_JINJA_CUSTOM_FILTERS,
             )
@@ -80,9 +86,7 @@ def compile_prompt_blocks(
             )
 
         elif block.block_type == "VARIABLE":
-            compiled_input: Optional[PromptRequestInput] = next(
-                (input_ for input_ in sanitized_inputs if input_.key == str(block.input_variable)), None
-            )
+            compiled_input = inputs_by_name.get(block.input_variable)
             if compiled_input is None:
                 raise PromptCompilationError(f"Input variable '{block.input_variable}' not found")
 
@@ -196,23 +200,26 @@ def _compile_chat_messages_as_prompt_blocks(chat_messages: list[ChatMessage]) ->
 
 def _compile_rich_text_block_as_value_block(
     block: RichTextPromptBlock,
-    inputs: list[PromptRequestInput],
+    inputs: list[PromptInput],
 ) -> CompiledValuePromptBlock:
     value: str = ""
+    inputs_by_name = {(input_.name if isinstance(input_, NumberInput) else input_.key): input_ for input_ in inputs}
     for child_block in block.blocks:
         if child_block.block_type == "PLAIN_TEXT":
             value += child_block.text
         elif child_block.block_type == "VARIABLE":
-            variable = next((input_ for input_ in inputs if input_.key == str(child_block.input_variable)), None)
-            if variable is None:
+            input = inputs_by_name.get(child_block.input_variable)
+            if input is None:
                 raise PromptCompilationError(f"Input variable '{child_block.input_variable}' not found")
-            elif variable.type == "STRING":
-                value += str(variable.value)
-            elif variable.type == "JSON":
-                value += json.dumps(variable.value, indent=4)
+            elif input.type == "STRING":
+                value += str(input.value)
+            elif input.type == "JSON":
+                value += json.dumps(input.value, indent=4)
+            elif input.type == "NUMBER":
+                value += str(input.value)
             else:
                 raise PromptCompilationError(
-                    f"Input variable '{child_block.input_variable}' must be of type STRING or JSON"
+                    f"Input variable '{child_block.input_variable}' has an invalid type: {input.type}"
                 )
         else:
             raise PromptCompilationError(f"Invalid child block_type for RICH_TEXT: {child_block.block_type}")
@@ -220,8 +227,8 @@ def _compile_rich_text_block_as_value_block(
     return CompiledValuePromptBlock(content=StringVellumValue(value=value), cache_config=block.cache_config)
 
 
-def _sanitize_inputs(inputs: list[PromptRequestInput]) -> list[PromptRequestInput]:
-    sanitized_inputs: list[PromptRequestInput] = []
+def _sanitize_inputs(inputs: Sequence[PromptInput]) -> list[PromptInput]:
+    sanitized_inputs: list[PromptInput] = []
     for input_ in inputs:
         if input_.type == "CHAT_HISTORY" and input_.value is None:
             sanitized_inputs.append(input_.model_copy(update={"value": cast(list[ChatMessage], [])}))

vellum/prompts/blocks/tests/test_compilation.py

@@ -10,6 +10,7 @@ from vellum import (
     VariablePromptBlock,
     VellumVariable,
 )
+from vellum.client.types.number_input import NumberInput
 from vellum.prompts.blocks.compilation import compile_prompt_blocks
 from vellum.prompts.blocks.types import CompiledChatMessagePromptBlock, CompiledValuePromptBlock
 
@@ -94,6 +95,33 @@ from vellum.prompts.blocks.types import CompiledChatMessagePromptBlock, Compiled
                 ),
             ],
         ),
+        (
+            [
+                ChatMessagePromptBlock(
+                    chat_role="USER",
+                    blocks=[
+                        RichTextPromptBlock(
+                            blocks=[
+                                PlainTextPromptBlock(text="Count to "),
+                                VariablePromptBlock(input_variable="count"),
+                            ]
+                        )
+                    ],
+                )
+            ],
+            [
+                # TODO: We don't yet have PromptRequestNumberInput. We should either add it or migrate
+                # Prompts to using these more generic inputs.
+                NumberInput(name="count", value=10),
+            ],
+            [VellumVariable(id="901ec2d6-430c-4341-b963-ca689006f5cc", type="NUMBER", key="count")],
+            [
+                CompiledChatMessagePromptBlock(
+                    role="USER",
+                    blocks=[CompiledValuePromptBlock(content=StringVellumValue(value="Count to 10.0"))],
+                ),
+            ],
+        ),
     ],
     ids=[
         "empty",
@@ -102,6 +130,7 @@ from vellum.prompts.blocks.types import CompiledChatMessagePromptBlock, Compiled
         "rich-text-no-variables",
         "rich-text-with-variables",
         "chat-message",
+        "number-input",
     ],
 )
 def test_compile_prompt_blocks__happy(blocks, inputs, input_variables, expected):

vellum/utils/templating/constants.py

@@ -10,7 +10,7 @@ import pydash
 import pytz
 import yaml
 
-from vellum.utils.templating.custom_filters import is_valid_json_string
+from vellum.utils.templating.custom_filters import is_valid_json_string, replace
 
 DEFAULT_JINJA_GLOBALS: Dict[str, Any] = {
     "datetime": datetime,
@@ -23,6 +23,10 @@ DEFAULT_JINJA_GLOBALS: Dict[str, Any] = {
     "re": re,
     "yaml": yaml,
 }
-DEFAULT_JINJA_CUSTOM_FILTERS: Dict[str, Callable[[Union[str, bytes]], bool]] = {
+
+FilterFunc = Union[Callable[[Union[str, bytes]], bool], Callable[[Any, Any, Any], str]]
+
+DEFAULT_JINJA_CUSTOM_FILTERS: Dict[str, FilterFunc] = {
     "is_valid_json_string": is_valid_json_string,
+    "replace": replace,
 }

vellum/utils/templating/custom_filters.py

@@ -1,5 +1,7 @@
 import json
-from typing import Union
+from typing import Any, Union
+
+from vellum.workflows.state.encoder import DefaultStateEncoder
 
 
 def is_valid_json_string(value: Union[str, bytes]) -> bool:
@@ -10,3 +12,22 @@ def is_valid_json_string(value: Union[str, bytes]) -> bool:
     except ValueError:
         return False
     return True
+
+
+def replace(s: Any, old: Any, new: Any) -> str:
+    def encode_to_str(obj: Any) -> str:
+        """Encode an object for template rendering using DefaultStateEncoder."""
+        try:
+            if isinstance(obj, str):
+                return obj
+            return json.dumps(obj, cls=DefaultStateEncoder)
+        except TypeError:
+            return str(obj)
+
+    if old == "":
+        return encode_to_str(s)
+
+    s_str = encode_to_str(s)
+    old_str = encode_to_str(old)
+    new_str = encode_to_str(new)
+    return s_str.replace(old_str, new_str)

vellum/utils/templating/render.py

@@ -1,8 +1,9 @@
 import json
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Dict, Optional
 
 from jinja2.sandbox import SandboxedEnvironment
 
+from vellum.utils.templating.constants import FilterFunc
 from vellum.utils.templating.exceptions import JinjaTemplateError
 from vellum.workflows.state.encoder import DefaultStateEncoder
 
@@ -18,7 +19,7 @@ def render_sandboxed_jinja_template(
     *,
     template: str,
     input_values: Dict[str, Any],
-    jinja_custom_filters: Optional[Dict[str, Callable[[Union[str, bytes]], bool]]] = None,
+    jinja_custom_filters: Optional[Dict[str, FilterFunc]] = None,
     jinja_globals: Optional[Dict[str, Any]] = None,
 ) -> str:
     """Render a Jinja template within a sandboxed environment."""
@@ -42,6 +43,8 @@ def render_sandboxed_jinja_template(
 
         rendered_template = jinja_template.render(input_values)
     except json.JSONDecodeError as e:
+        if not e.doc:
+            raise JinjaTemplateError("Unable to render jinja template:\n" "Cannot run json.loads() on empty input")
         if e.msg == "Invalid control character at":
             raise JinjaTemplateError(
                 "Unable to render jinja template:\n"

vellum/utils/templating/tests/test_custom_filters.py

@@ -0,0 +1,19 @@
+import pytest
+
+from vellum.utils.templating.custom_filters import replace
+
+
+@pytest.mark.parametrize(
+    ["input_str", "old", "new", "expected"],
+    [
+        ("foo", "foo", "bar", "bar"),
+        ({"message": "hello"}, "hello", "world", '{"message": "world"}'),
+        ("Value: 123", 123, 456, "Value: 456"),
+        (123, 2, 4, "143"),
+        ("", "", "", ""),
+        ("foo", "", "bar", "foo"),
+    ],
+)
+def test_replace(input_str, old, new, expected):
+    actual = replace(input_str, old, new)
+    assert actual == expected

vellum/workflows/errors/types.py

@@ -13,6 +13,7 @@ class WorkflowErrorCode(Enum):
     INVALID_INPUTS = "INVALID_INPUTS"
     INVALID_OUTPUTS = "INVALID_OUTPUTS"
     INVALID_STATE = "INVALID_STATE"
+    INVALID_CODE = "INVALID_CODE"
     INVALID_TEMPLATE = "INVALID_TEMPLATE"
     INTERNAL_ERROR = "INTERNAL_ERROR"
     NODE_EXECUTION = "NODE_EXECUTION"
@@ -54,6 +55,7 @@ _WORKFLOW_EVENT_ERROR_CODE_TO_WORKFLOW_ERROR_CODE: Dict[WorkflowExecutionEventEr
     "INTERNAL_SERVER_ERROR": WorkflowErrorCode.INTERNAL_ERROR,
     "NODE_EXECUTION": WorkflowErrorCode.NODE_EXECUTION,
     "LLM_PROVIDER": WorkflowErrorCode.PROVIDER_ERROR,
+    "INVALID_CODE": WorkflowErrorCode.INVALID_CODE,
     "INVALID_TEMPLATE": WorkflowErrorCode.INVALID_TEMPLATE,
     "USER_DEFINED_ERROR": WorkflowErrorCode.USER_DEFINED_ERROR,
 }
@@ -71,6 +73,7 @@ _WORKFLOW_ERROR_CODE_TO_VELLUM_ERROR_CODE: Dict[WorkflowErrorCode, VellumErrorCo
     WorkflowErrorCode.INVALID_INPUTS: "INVALID_INPUTS",
     WorkflowErrorCode.INVALID_OUTPUTS: "INVALID_REQUEST",
     WorkflowErrorCode.INVALID_STATE: "INVALID_REQUEST",
+    WorkflowErrorCode.INVALID_CODE: "INVALID_CODE",
     WorkflowErrorCode.INVALID_TEMPLATE: "INVALID_INPUTS",
     WorkflowErrorCode.INTERNAL_ERROR: "INTERNAL_SERVER_ERROR",
     WorkflowErrorCode.NODE_EXECUTION: "USER_DEFINED_ERROR",

vellum/workflows/nodes/core/templating_node/node.py

@@ -1,6 +1,6 @@
-from typing import Any, Callable, ClassVar, Dict, Generic, Mapping, Tuple, Type, TypeVar, Union, get_args
+from typing import Any, ClassVar, Dict, Generic, Mapping, Tuple, Type, TypeVar, get_args
 
-from vellum.utils.templating.constants import DEFAULT_JINJA_CUSTOM_FILTERS, DEFAULT_JINJA_GLOBALS
+from vellum.utils.templating.constants import DEFAULT_JINJA_CUSTOM_FILTERS, DEFAULT_JINJA_GLOBALS, FilterFunc
 from vellum.utils.templating.exceptions import JinjaTemplateError
 from vellum.utils.templating.render import render_sandboxed_jinja_template
 from vellum.workflows.errors import WorkflowErrorCode
@@ -54,7 +54,7 @@ class TemplatingNode(BaseNode[StateType], Generic[StateType, _OutputType], metac
     inputs: ClassVar[EntityInputsInterface]
 
     jinja_globals: Dict[str, Any] = DEFAULT_JINJA_GLOBALS
-    jinja_custom_filters: Mapping[str, Callable[[Union[str, bytes]], bool]] = DEFAULT_JINJA_CUSTOM_FILTERS
+    jinja_custom_filters: Mapping[str, FilterFunc] = DEFAULT_JINJA_CUSTOM_FILTERS
 
     class Outputs(BaseNode.Outputs):
         """

vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py

@@ -1,8 +1,11 @@
+import pytest
 import json
-from typing import List
+from typing import List, Union
 
 from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.function_call import FunctionCall
+from vellum.client.types.function_call_vellum_value import FunctionCallVellumValue
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.state import BaseState
@@ -155,3 +158,65 @@ def test_templating_node__function_call_output():
 
     # THEN the output is the expected function call
     assert outputs.result == FunctionCall(name="test", arguments={"key": "value"})
+
+
+def test_templating_node__blank_json_input():
+    """Test that templating node properly handles blank JSON input."""
+
+    # GIVEN a templating node that tries to parse blank JSON
+    class BlankJsonTemplateNode(TemplatingNode[BaseState, Json]):
+        template = "{{ json.loads(data) }}"
+        inputs = {
+            "data": "",  # Blank input
+        }
+
+    # WHEN the node is run
+    node = BlankJsonTemplateNode()
+
+    # THEN it should raise an appropriate error
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert "Unable to render jinja template:\nCannot run json.loads() on empty input" in str(exc_info.value)
+
+
+def test_templating_node__union_float_int_output():
+    # GIVEN a templating node that outputs either a float or an int
+    class UnionTemplateNode(TemplatingNode[BaseState, Union[float, int]]):
+        template = """{{ obj[\"score\"] | float }}"""
+        inputs = {
+            "obj": {"score": 42.5},
+        }
+
+    # WHEN the node is run
+    node = UnionTemplateNode()
+    outputs = node.run()
+
+    # THEN it should correctly parse as a float
+    assert outputs.result == 42.5
+
+
+def test_templating_node__replace_filter():
+    # GIVEN a templating node that outputs a complex object
+    class ReplaceFilterTemplateNode(TemplatingNode[BaseState, Json]):
+        template = """{{- prompt_outputs | selectattr(\'type\', \'equalto\', \'FUNCTION_CALL\') \
+        | list | replace(\"\\n\",\",\") -}}"""
+        inputs = {
+            "prompt_outputs": [FunctionCallVellumValue(value=FunctionCall(name="test", arguments={"key": "value"}))]
+        }
+
+    # WHEN the node is run
+    node = ReplaceFilterTemplateNode()
+    outputs = node.run()
+
+    # THEN the output is the expected JSON
+    assert outputs.result == [
+        {
+            "type": "FUNCTION_CALL",
+            "value": {
+                "name": "test",
+                "arguments": {"key": "value"},
+                "id": None,
+            },
+        }
+    ]

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

@@ -93,13 +93,14 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
         log: str
 
     def run(self) -> Outputs:
-        input_values = self._compile_code_inputs()
         output_type = self.__class__.get_output_type()
         code = self._resolve_code()
         if not self.packages and self.runtime == "PYTHON_3_11_6":
-            logs, result = run_code_inline(code, input_values, output_type)
+            logs, result = run_code_inline(code, self.code_inputs, output_type)
             return self.Outputs(result=result, log=logs)
+
         else:
+            input_values = self._compile_code_inputs()
             expected_output_type = primitive_type_to_vellum_variable_type(output_type)
 
             code_execution_result = self._context.vellum_client.execute_code(
@@ -131,7 +132,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                 compiled_inputs.append(
                     StringInput(
                         name=input_name,
-                        value=str(input_value),
+                        value=input_value,
                     )
                 )
             elif isinstance(input_value, VellumSecret):
@@ -193,7 +194,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                 )
             else:
                 raise NodeException(
-                    message=f"Unrecognized input type for input '{input_name}'",
+                    message=f"Unrecognized input type for input '{input_name}': {input_value.__class__.__name__}",
                     code=WorkflowErrorCode.INVALID_INPUTS,
                 )
 

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -7,6 +7,7 @@ from vellum.client.types.code_execution_package import CodeExecutionPackage
 from vellum.client.types.code_executor_secret_input import CodeExecutorSecretInput
 from vellum.client.types.function_call import FunctionCall
 from vellum.client.types.number_input import NumberInput
+from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
@@ -559,3 +560,53 @@ def main(arg1: List[Dict]) -> float:
 
     # AND we should not have invoked the Code via Vellum since it's running inline
     vellum_client.execute_code.assert_not_called()
+
+
+def test_run_node__code_execution_error():
+    # GIVEN a node that will raise an error during execution
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        code = """\
+def main(arg1: int, arg2: int) -> int:
+    return arg1 + arg2 + arg3
+"""
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "arg1": 1,
+            "arg2": 2,
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+
+    # THEN it should raise a NodeException with the execution error
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # AND the error should contain the execution error details
+    assert "name 'arg3' is not defined" in str(exc_info.value)
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_CODE
+
+
+def test_run_node__array_of_bools_input():
+    # GIVEN a node that will raise an error during execution
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
+        code = """\
+def main(arg1: list[bool]) -> int:
+    return len(arg1)
+"""
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "arg1": [True, False, True],
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+
+    # THEN it should raise a NodeException with the execution error
+    outputs = node.run()
+
+    # AND the error should contain the execution error details
+    assert outputs == {"result": 3, "log": ""}

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

@@ -1,14 +1,13 @@
 import io
 import os
 import re
-from typing import Any, List, Tuple, Union, get_args, get_origin
+from typing import Any, Tuple, Union, get_args, get_origin
 
 from pydantic import BaseModel, ValidationError
 
-from vellum import VellumValue
-from vellum.client.types.code_executor_input import CodeExecutorInput
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.types.core import EntityInputsInterface
 
 
 def read_file_from_path(node_filepath: str, script_filepath: str) -> Union[str, None]:
@@ -70,13 +69,11 @@ def _clean_for_dict_wrapper(obj):
 
 def run_code_inline(
     code: str,
-    input_values: List[CodeExecutorInput],
+    inputs: EntityInputsInterface,
     output_type: Any,
 ) -> Tuple[str, Any]:
     log_buffer = io.StringIO()
 
-    VELLUM_TYPES = get_args(VellumValue)
-
     def wrap_value(value):
         if isinstance(value, list):
             return ListWrapper(
@@ -84,7 +81,7 @@ def run_code_inline(
                     # Convert VellumValue to dict with its fields
                     (
                         item.model_dump()
-                        if isinstance(item, VELLUM_TYPES)
+                        if isinstance(item, BaseModel)
                         else _clean_for_dict_wrapper(item) if isinstance(item, (dict, list)) else item
                     )
                     for item in value
@@ -93,18 +90,23 @@ def run_code_inline(
         return _clean_for_dict_wrapper(value)
 
     exec_globals = {
-        "__arg__inputs": {input_value.name: wrap_value(input_value.value) for input_value in input_values},
+        "__arg__inputs": {name: wrap_value(value) for name, value in inputs.items()},
         "__arg__out": None,
         "print": lambda *args, **kwargs: log_buffer.write(f"{' '.join(args)}\n"),
     }
-    run_args = [f"{input_value.name}=__arg__inputs['{input_value.name}']" for input_value in input_values]
+    run_args = [f"{name}=__arg__inputs['{name}']" for name in inputs.keys()]
     execution_code = f"""\
 {code}
 
 __arg__out = main({", ".join(run_args)})
 """
-
-    exec(execution_code, exec_globals)
+    try:
+        exec(execution_code, exec_globals)
+    except Exception as e:
+        raise NodeException(
+            code=WorkflowErrorCode.INVALID_CODE,
+            message=str(e),
+        )
 
     logs = log_buffer.getvalue()
     result = exec_globals["__arg__out"]

vellum/workflows/nodes/utils.py

@@ -2,7 +2,7 @@ from functools import cache
 import json
 import sys
 from types import ModuleType
-from typing import Any, Callable, Optional, Type, TypeVar, get_args, get_origin
+from typing import Any, Callable, Optional, Type, TypeVar, Union, get_args, get_origin
 
 from pydantic import BaseModel
 
@@ -113,6 +113,14 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
         except json.JSONDecodeError:
             raise ValueError("Invalid JSON format for result_as_str")
 
+    if get_origin(output_type) is Union:
+        for inner_type in get_args(output_type):
+            try:
+                return parse_type_from_str(result_as_str, inner_type)
+            except ValueError:
+                continue
+        raise ValueError(f"Could not parse with any of the Union types: {output_type}")
+
     if issubclass(output_type, BaseModel):
         try:
             data = json.loads(result_as_str)

vellum/workflows/state/base.py

@@ -18,7 +18,13 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.types.stack import Stack
-from vellum.workflows.types.utils import datetime_now, deepcopy_with_exclusions, get_class_by_qualname, infer_types
+from vellum.workflows.types.utils import (
+    datetime_now,
+    deepcopy_with_exclusions,
+    get_class_attr_names,
+    get_class_by_qualname,
+    infer_types,
+)
 
 if TYPE_CHECKING:
     from vellum.workflows.nodes.bases import BaseNode
@@ -47,6 +53,20 @@ class _BaseStateMeta(type):
 
         return super().__getattribute__(name)
 
+    def __iter__(cls) -> Iterator[StateValueReference]:
+        # We iterate through the inheritance hierarchy to find all the StateValueReference attached to this
+        # Inputs class. __mro__ is the method resolution order, which is the order in which base classes are resolved.
+        for resolved_cls in cls.__mro__:
+            attr_names = get_class_attr_names(resolved_cls)
+            for attr_name in attr_names:
+                if attr_name == "meta":
+                    continue
+                attr_value = getattr(resolved_cls, attr_name)
+                if not isinstance(attr_value, (StateValueReference)):
+                    continue
+
+                yield attr_value
+
 
 class _SnapshottableDict(dict, _Snapshottable):
     def __setitem__(self, key: Any, value: Any) -> None:

{vellum_ai-0.14.0.dist-info → vellum_ai-0.14.2.dist-info}/METADATA

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