vellum-ai 0.14.0__py3-none-any.whl → 0.14.1__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.1",
         }
         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/render.py

@@ -42,6 +42,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/workflows/nodes/core/templating_node/tests/test_templating_node.py

@@ -1,8 +1,10 @@
+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.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 +157,39 @@ 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

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.1.dist-info}/METADATA

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

{vellum_ai-0.14.0.dist-info → vellum_ai-0.14.1.dist-info}/RECORD

@@ -19,9 +19,9 @@ vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/base.py,sha256=3ZFUYRNKL24fBqXhKpa_Dq2W1a-a86J20dmJYA3H2eY,1755
+vellum_ee/workflows/display/base.py,sha256=i8o7Ex0S9BVub_91s6D2a8kHGcv3ro5XGEwSzeVSz-Y,2083
 vellum_ee/workflows/display/nodes/__init__.py,sha256=5XOcZJXYUgaLS55QgRJzyQ_W1tpeprjnYAeYVezqoGw,160
-vellum_ee/workflows/display/nodes/base_node_display.py,sha256=j5AoStKNjSmSdxKNqJsi_de4Lhi1WmK0IjdXLQFf2oE,16857
+vellum_ee/workflows/display/nodes/base_node_display.py,sha256=En8Ir2e1mpExkAi1T8ExOEpZbv5gu6OetaNq35-_WSY,16317
 vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=pLO0dORfRu--Ne9NgoyFT_CNjfpr5fGCsgbsMkUF5GM,2845
 vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=0S6ksPp53WXWh1RQVH71nj2bkCWBj4ZaFYhTxW3N2F4,1230
 vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -47,46 +47,49 @@ vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=6
 vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=JVIMPR3WpveOCWZubHKZkE04mavnTdb_9QY_r3XliRg,3424
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py,sha256=ulrpoYUW-5kIxfG4Lf5F2p0k_EoYKhmahEbF3P_eruM,1648
+vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py,sha256=bg9INsXiWfyK047u8TD1oEOFYrqDq8GC7Hvgz69n7BE,1988
 vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py,sha256=mtzB8LJlFCHVFM4H5AanLp29gQfaVmnN4A4iaRGJHoI,2427
 vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=3uT7Gbc0f_mQ3u8uZuCWd0mJ4GtWbz2gbUMySYaVlNE,3774
 vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=EoU1J7HfcszAZr7ROy_xsNhaDBRiI95-wTK-OzgBwvg,5861
 vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=OEGHjQSbuUgJexXI1aubYW33z2F_YdkhQ8REahfz864,4320
 vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=h4bE187MyM0eT_elFnOxwsCeMHeuzSK0orjgAy8kV1k,4952
+vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=1EEvkKQRfOKlnpLxE9-hKSsVLLaelM39LY7007LM5dg,4983
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=EenmEdBtHUFQ0OS-kE7Vboax3JnDdj-K4Qixt5oR0Po,2253
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=A1-tIpC5KIKG9JA_rkd1nLS8zUG3Kb4QiVdvb3boFxE,2509
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=KYdohS5pRgHM0DcUaK0tHa40f0gSvDKi2K5On0zNEU8,8305
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=1cszL6N6FNGVm61MOa7AEiHnF0QjZWqDQuPOp4yiG94,18277
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=-12ZkZb3f5gyoNASV2yeQtMo5HmNsVEo8nXwL6IC-I8,6261
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=6th6kCwzql6lddjkTQx4Jbvvs4ChqtJwctW-B4QuBhI,37352
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py,sha256=EbVgg_3_ipTt3MOop4RARX0fmNjwqZtkhIXzx9nGw7Y,4487
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py,sha256=5LQ36vImbVRmLpOdhWKtGmr9pypeZc0HAzIxwQoKtWA,15890
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=3lmg0bP6QwnHjQCpxz3IiucowmgAcjelHRTE9JRQ9-A,29152
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py,sha256=ddk2tx7dFqQ7o2JtRZz2Mv6GHEwjJSF_u-1Zg83q7to,47564
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=02gG5C0wHbjzuTgVYCsNrtW1kEOeaM_5zXmznaUKzgk,6079
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py,sha256=Ah8CxAxAQYnxsNasLB8esN_c9eRDQnmV_aDWC9Kp98s,5746
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py,sha256=0xK9TFqPD5Hy65T-iacZ8pjYHD8XkzZXS35rLwof7uc,7427
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py,sha256=IRazH2QR6F8RGqNemEnHueyj5DtEa6rFTYhT16S4jI8,15917
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=PiHqgj-vLhlaOHoQFVVEW2YydYWLI0mX6nq_sbzZiy4,29233
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py,sha256=R8DW1DUb0DOSLtnF2E1HaCTmtpG-ski0LfcM2WeLVNo,47672
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_default_state_serialization.py,sha256=z9FSufJeV-003R87wi_Lx4mZewdeeOeTPCGNc9vg8vY,8819
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=wgedEa2IVP2ssH_FLghoEmLgpJR41AY-iNIw1SESeqA,6106
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py,sha256=yIq5UexHAk7Afh_6sMlXNAUeKdckBUvCBVIEDLEnhFs,5773
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py,sha256=0N7lqB2WMvSItUY78EWpF-g4fXvhkc4NCY-Ws57OYsQ,7446
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py,sha256=XRtXTWtNXHe1g8E_mNcAhYZVjc0UcGlJRjq34_sd0jo,787
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py,sha256=HEfRYBVwZ6fy0hHhgsgTBEQJJAp_0wsaCJ_OtpZGdqE,20578
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py,sha256=L1SrD5KfZjlAVh2BkOe3pCQOstVf_SiqD0KGoZlknU8,16145
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py,sha256=ojzdgGb2zGYloOX57V_FyXxjUqOnWMN6z8GUlF8yDMA,8422
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py,sha256=SuADagoAr-IWZN5az98axWnXcSrlG6gXV72Q-ThozVU,8622
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py,sha256=NcpcC9vENF0Z3Vg7b0G0WV7EJScmBKg-JtjQENMKvn0,12989
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py,sha256=eQNyiNQG7Cgx0QW_xlN9tQql4D2ak_P2WB-OgMkAVEs,11295
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py,sha256=9_mC6-_u3qxJZzkrYrDtzWhpKlnlOD4oXSj8Ztnwet4,7791
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py,sha256=1HaEwIbxzWl2OEEVlsNcHFQKMhUKQPlsIMaHnjvMODI,3754
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=5eKEMwrKT5O8dPsOgQ89iZwzpA304jneimQeT9GlpLk,2428
-vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=xmKyOKqGRhWUXXbTbCAXX40qINcvK7mjpEfXz0Ny21s,7404
-vellum_ee/workflows/display/types.py,sha256=s1w2KELKI2kde4G2M3iniOEdPObNKYgYtr31sAgZqRI,2644
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py,sha256=cKM87WUrBZwVauP7VD4rQdObd2jllHJJRHWDGL1tgMI,20605
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py,sha256=PbIezAnkGubhY_e5dpOaaaVZ9VeoYJKqvCKbIbjucfU,16172
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py,sha256=ObCnoIVJM434jJBC3_kOTKbY4sCDHlknwMTFWRGiOtU,8449
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py,sha256=-bGp82C5_XgcVxeM80zWL9viogxvE4DGk7-5H__wwxA,8649
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py,sha256=elc4ak2OyI4y5P6GPnK2RwuZLcx4FL6hHlmhY3ARk6k,13008
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py,sha256=BzFNl9ECeGh0krm-CUjbBQQq0g7krANsp0Sh-j5dAkc,11322
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py,sha256=xXtW915v9yxxKlyu5ObzKHyJYMvobvev3ermX61SGx4,7818
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py,sha256=NdhE3lm7RMQ8DqkraPSq24IbOxNla9unbs4tsMWRzm4,3781
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=gHKU8Vg960ooV3uXqM2LMnVS-mGbv3aagGozQuTVTjI,2455
+vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=huKAOeMJ2MKmp6XtbvMJTUadqynoV40Ypoz9jsBEBEQ,7431
+vellum_ee/workflows/display/types.py,sha256=bm355ZcZ-7WaesBvIB8RbWDD7qiFmazfcNb6tEZWx3M,2940
 vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/utils/vellum.py,sha256=IGHobHz0QG8TWdJNxlQ81khCmJv7piLjE64l4zc33kI,7743
-vellum_ee/workflows/display/vellum.py,sha256=1mcZQQ4gd5GYrVKgIQ3ejy8SqRss2CWUL0KMf44zgUc,8122
+vellum_ee/workflows/display/utils/expressions.py,sha256=9FpOslDI-RCR5m4TgAu9KCHh4aTVnh7CHR2ykyMUDw0,1151
+vellum_ee/workflows/display/utils/vellum.py,sha256=UjK_RxnSEmlIu9klGCPWU5RAQBmgZ7cRbRdgxaTbubE,8081
+vellum_ee/workflows/display/vellum.py,sha256=PVte2gcp1_4g-8ZaM0XVAaDwf3O-kSHy7-5Q-Co-YG0,8463
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=8inuxnxeJbysEhJIGdxLZTvQ2aPJj-OSGmSeotwe7fo,15782
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=Amow11LhE-GL3j8bNZ8rK1iYKjkNOhH1ECMQ-dvyLAU,17204
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
 vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=yl1ytpl9_lijOGeDPWSypWYRJ7aOEVA7NgUg81jTuCs,2229
-vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=IiHfMglir81cwPFYnwvgynp79cgzFyYVkMJAABIi3BQ,17035
+vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=12miGvHL9_ZzXaU5r1u1TVTGSBrs8LHBVJU-DKfS0JE,18802
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
 vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -116,7 +119,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=j6zi0NZ4BMC6JrwckvzMWuG5x8KoOvO4KqsLhvVCa68,117624
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=16qoykymtSsNTtgTE4q07vaJfALO7E-GbQ6P8GVl1Mw,1868
+vellum/client/core/client_wrapper.py,sha256=AN4UTP8OvOr6qMSP_VqGlHKGOVsUnQj3CLqFDhB3p5E,1868
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -706,10 +709,10 @@ vellum/plugins/utils.py,sha256=U9ZY9KdE3RRvbcG01hXxu9CvfJD6Fo7nJDgcHjQn0FI,606
 vellum/plugins/vellum_mypy.py,sha256=QTuMSq6PiZW1dyTUZ5Bf1d4XkgFj0TKAgZLP8f4UgL4,27914
 vellum/prompts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/prompts/blocks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/prompts/blocks/compilation.py,sha256=ISuvDHaeVCPb1L7l4umORCECkDn0-rvE49hopz6c2gM,9222
+vellum/prompts/blocks/compilation.py,sha256=P2gdL4bRR1j0IbvQFzS13cJzGpnaDZFhdV3x5COVvs0,9481
 vellum/prompts/blocks/exceptions.py,sha256=vmk5PV6Vyw9nKjZYQDUDW0LH8MfQNIgFvFb_mFWdIRI,50
 vellum/prompts/blocks/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/prompts/blocks/tests/test_compilation.py,sha256=0DhMoc4huHR6YnNL-0aBLmWSyUfw2BpRq_gEdKsQmAc,3693
+vellum/prompts/blocks/tests/test_compilation.py,sha256=WhR8L7AzxssZiizItAbbtQMXZ3An5RMUws0kfCwJBgw,4841
 vellum/prompts/blocks/types.py,sha256=6aSJQco-5kKeadfKVVXF_SrQPlIJgMYVNc-C7so1sY8,975
 vellum/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/resources/__init__.py,sha256=sQWK7g_Z4EM7pa7fy6vy3d_DMdTJ4wVcozBn3Lx4Qpo,141
@@ -1259,7 +1262,7 @@ vellum/utils/templating/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG
 vellum/utils/templating/constants.py,sha256=XTNmDsKa7byjw4GMZmzx2dUeYUTeMLZrPgRHcc80Kvc,613
 vellum/utils/templating/custom_filters.py,sha256=Q0DahYRHP4KfaUXDt9XxN-DFLBrAxlv90yaVqxScoUw,264
 vellum/utils/templating/exceptions.py,sha256=cDp140PP4OnInW4qAvg3KqiSiF70C71UyEAKRBR1Abo,46
-vellum/utils/templating/render.py,sha256=OwdZTJlQv_qsygMX3LOFr4d1_2oz3r9vGKyj7cX05VE,1876
+vellum/utils/templating/render.py,sha256=2n6M10p2WtM2_6I10PKOIyH_gTmAo5EkxDYL2DR7F4U,2014
 vellum/utils/typing.py,sha256=wx_daFqD69cYkuJTVnvNrpjhqC3uuhbnyJ9_bIwC9OU,327
 vellum/utils/uuid.py,sha256=Ch6wWRgwICxLxJCTl5iE3EdRlZj2zADR-zUMUtjcMWM,214
 vellum/version.py,sha256=jq-1PlAYxN9AXuaZqbYk9ak27SgE2lw9Ia5gx1b1gVI,76
@@ -1346,7 +1349,7 @@ vellum/workflows/nodes/core/retry_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TI
 vellum/workflows/nodes/core/retry_node/tests/test_node.py,sha256=RM_OHwxrHwyxvlQQBJPqVBxpedFuWQ9h2-Xa3kP75sc,4399
 vellum/workflows/nodes/core/templating_node/__init__.py,sha256=GmyuYo81_A1_Bz6id69ozVFS6FKiuDsZTiA3I6MaL2U,70
 vellum/workflows/nodes/core/templating_node/node.py,sha256=zCYhq88qLTvoC9LetVrD9sLXkwHZsaWekxMhru_nV70,3752
-vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=5iZWQWdJKDHMXBY8bhpb-Dpy9FTfW1HXxGUTivykZAA,4621
+vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=j6v3TTD4Fjq-sWcuOZOzX-okFUC18q2VKiPT9ot1bk8,5830
 vellum/workflows/nodes/core/try_node/__init__.py,sha256=JVD4DrldTIqFQQFrubs9KtWCCc0YCAc7Fzol5ZWIWeM,56
 vellum/workflows/nodes/core/try_node/node.py,sha256=bk2uhYUl10yaPJlOBWxiL7igTUrL_7mM9S2nvsdWB68,4242
 vellum/workflows/nodes/core/try_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1412,7 +1415,7 @@ vellum/workflows/nodes/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQ
 vellum/workflows/nodes/experimental/openai_chat_completion_node/__init__.py,sha256=lsyD9laR9p7kx5-BXGH2gUTM242UhKy8SMV0SR6S2iE,90
 vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py,sha256=1EGeiaT-Zoo6pttQFKKBcdf3dmhAbjKGaErYD5FFwlc,10185
 vellum/workflows/nodes/mocks.py,sha256=gvM2tyoe-V84jFbFdhQsyGAPyQBzmjn_CkhT_yxccgY,499
-vellum/workflows/nodes/utils.py,sha256=T7krLipDSI007JkkH2zsfyQ-tfOBH4hyFzQ0YJKYpuw,4025
+vellum/workflows/nodes/utils.py,sha256=YRj3qIz6--N6CuDiKalsWpBmUR6z7WcRhqtLxoenDZE,4354
 vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
 vellum/workflows/outputs/base.py,sha256=RPlkKwVml6Ln8-K7kcpRl8-T-0OmB6css_UaKJdjAA8,8581
 vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
@@ -1437,7 +1440,7 @@ vellum/workflows/runner/__init__.py,sha256=i1iG5sAhtpdsrlvwgH6B-m49JsINkiWyPWs8v
 vellum/workflows/runner/runner.py,sha256=eY4axSyxotT_NHB3CYFA1hDSoM-eY4k3Gq5NEe7dpQw,29208
 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=O8BZpXTGsMVEltTs70iIms2PtZ3Qd0imf_OyoB1pKaI,14775
+vellum/workflows/state/base.py,sha256=Vkhneko3VlQrPsMLU1PYSzXU_W1u7_AraJsghiv5O-4,15512
 vellum/workflows/state/context.py,sha256=80b7j11AuG6engdur-AH0dUPnbwu8QyzF_-crWmvUvI,2588
 vellum/workflows/state/encoder.py,sha256=WdUidpOaBDx5lilJl8V8McFDHQYiCLCJR9dmktdzdZY,1836
 vellum/workflows/state/store.py,sha256=VYGBQgN1bpd1as5eGiouV_7scg8QsRs4_1aqZAGIsRQ,793
@@ -1468,8 +1471,8 @@ vellum/workflows/workflows/base.py,sha256=tBivLLLPtMyMi2Bd5WjbwQ2_zZFpG8rNfEidfQ
 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=dz5F3DuGOk9qUMjmNd-GdEE3320G5ko5nJ6J0QJyVcY,2659
-vellum_ai-0.14.0.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.0.dist-info/METADATA,sha256=whh5Vpg1kDk2eaO-spcwQdR3inmlyl2YWrPoGWgVbYg,5407
-vellum_ai-0.14.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.0.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.0.dist-info/RECORD,,
+vellum_ai-0.14.1.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.1.dist-info/METADATA,sha256=TGCti5BkhYzjafUugASpEkqaT-cfB5-0K0zWZQPi3II,5407
+vellum_ai-0.14.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.1.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.1.dist-info/RECORD,,

vellum_ee/workflows/display/base.py

@@ -31,6 +31,20 @@ WorkflowInputsDisplayType = TypeVar("WorkflowInputsDisplayType", bound=WorkflowI
 WorkflowInputsDisplayOverridesType = TypeVar("WorkflowInputsDisplayOverridesType", bound=WorkflowInputsDisplayOverrides)
 
 
+@dataclass
+class StateValueDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class StateValueDisplay(StateValueDisplayOverrides):
+    pass
+
+
+StateValueDisplayType = TypeVar("StateValueDisplayType", bound=StateValueDisplay)
+StateValueDisplayOverridesType = TypeVar("StateValueDisplayOverridesType", bound=StateValueDisplayOverrides)
+
+
 @dataclass
 class EdgeDisplayOverrides:
     id: UUID

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -34,6 +34,7 @@ from vellum.workflows.references import OutputReference
 from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.references.state_value import StateValueReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
@@ -44,6 +45,7 @@ from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator, primitive_to_vellum_value
 from vellum_ee.workflows.display.vellum import CodeResourceDefinition, GenericNodeDisplayData
 
@@ -343,28 +345,8 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             return self.serialize_value(display_context, value._value)
 
         if isinstance(value, LazyReference):
-            if isinstance(value._get, str):
-                reference_parts = value._get.split(".")
-                if len(reference_parts) < 3:
-                    raise Exception(
-                        f"Failed to parse lazy reference: {value._get}. Only Node Output references are supported."
-                    )
-
-                output_name = reference_parts[-1]
-                nested_class_name = reference_parts[-2]
-                if nested_class_name != "Outputs":
-                    raise Exception(
-                        f"Failed to parse lazy reference: {value._get}. Outputs are the only node reference supported."
-                    )
-
-                node_class_name = ".".join(reference_parts[:-2])
-                for node in display_context.node_displays.keys():
-                    if node.__name__ == node_class_name:
-                        return self.serialize_value(display_context, getattr(node.Outputs, output_name))
-
-                raise NotImplementedError(f"Failed to find a LazyReference for: {value._get}")
-
-            return self.serialize_value(display_context, value._get())
+            child_descriptor = get_child_descriptor(value, display_context)
+            return self.serialize_value(display_context, child_descriptor)
 
         if isinstance(value, WorkflowInputReference):
             workflow_input_display = display_context.global_workflow_input_displays[value]
@@ -373,6 +355,13 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 "input_variable_id": str(workflow_input_display.id),
             }
 
+        if isinstance(value, StateValueReference):
+            state_value_display = display_context.global_state_value_displays[value]
+            return {
+                "type": "STATE_VALUE",
+                "state_variable_id": str(state_value_display.id),
+            }
+
         if isinstance(value, OutputReference):
             upstream_node, output_display = display_context.global_node_output_displays[value]
             upstream_node_display = display_context.global_node_displays[upstream_node]

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

@@ -0,0 +1,52 @@
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.references.lazy import LazyReference
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+
+
+def test_serialize_node__lazy_reference_in_prompt_inputs():
+    # GIVEN a prompt node with a lazy reference in the prompt inputs
+    class LazyReferencePromptNode(InlinePromptNode):
+        prompt_inputs = {"attr": LazyReference[str]("OtherNode.Outputs.result")}
+        blocks = []
+        ml_model = "gpt-4o"
+
+    class OtherNode(BaseNode):
+        class Outputs:
+            result: str
+
+    # AND a workflow with both nodes
+    class Workflow(BaseWorkflow):
+        graph = LazyReferencePromptNode >> OtherNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the attribute reference
+    lazy_reference_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(LazyReferencePromptNode.__id__)
+    )
+
+    assert lazy_reference_node["inputs"] == [
+        {
+            "id": "fba6a4d5-835a-4e99-afb7-f6a4aed15110",
+            "key": "attr",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "NODE_OUTPUT",
+                        "data": {
+                            "node_id": str(OtherNode.__id__),
+                            "output_id": "7f377cb8-4eca-4f1c-9239-9925f9495d84",
+                        },
+                    }
+                ],
+            },
+        }
+    ]

vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -25,6 +25,7 @@ def test_vellum_workflow_display__serialize_empty_workflow():
     # THEN it should return the expected config
     assert exec_config == {
         "input_variables": [],
+        "state_variables": [],
         "output_variables": [],
         "workflow_raw_data": {
             "definition": {

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py

@@ -4,10 +4,11 @@ from typing import Any, Dict, Tuple, Type
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.references.output import OutputReference
+from vellum.workflows.references.state_value import StateValueReference
 from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
-from vellum_ee.workflows.display.base import WorkflowInputsDisplayType
+from vellum_ee.workflows.display.base import StateValueDisplayType, WorkflowInputsDisplayType
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
@@ -22,6 +23,7 @@ def serialize_node():
         node_class: Type[NodeType],
         base_class: type[BaseNodeDisplay[Any]] = BaseNodeDisplay,
         global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {},
+        global_state_value_displays: Dict[StateValueReference, StateValueDisplayType] = {},
         global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = {},
         global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = {},
     ) -> JsonObject:
@@ -37,6 +39,7 @@ def serialize_node():
             ),
             node_displays={node_class: node_display},
             global_workflow_input_displays=global_workflow_input_displays,
+            global_state_value_displays=global_state_value_displays,
             global_node_displays=global_node_displays,
             global_node_output_displays=global_node_output_displays,
         )