vellum-ai 0.12.16__py3-none-any.whl → 0.12.17__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.12.16",
+            "X-Fern-SDK-Version": "0.12.17",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/plugins/vellum_mypy.py

@@ -4,6 +4,7 @@ from typing import Callable, Dict, List, Optional, Set, Type
 from mypy.nodes import (
     AssignmentStmt,
     CallExpr,
+    ClassDef,
     Decorator,
     MemberExpr,
     NameExpr,
@@ -198,13 +199,27 @@ class VellumMypyPlugin(Plugin):
             return
 
         current_node_outputs = node_info.names.get("Outputs")
-        if not current_node_outputs:
+        if not current_node_outputs and isinstance(base_node_outputs.node, TypeInfo):
             node_info.names["Outputs"] = base_node_outputs.copy()
-            new_outputs_sym = node_info.names["Outputs"].node
-            if isinstance(new_outputs_sym, TypeInfo):
-                result_sym = new_outputs_sym.names[attribute_name].node
-                if isinstance(result_sym, Var):
-                    result_sym.type = base_node_resolved_type
+
+            new_outputs_fullname = f"{node_info.fullname}.Outputs"
+            new_outputs_defn_raw = base_node_outputs.node.defn.serialize()
+            new_outputs_defn_raw["fullname"] = new_outputs_fullname
+            new_outputs_defn = ClassDef.deserialize(new_outputs_defn_raw)
+            new_outputs_sym = TypeInfo(
+                names=base_node_outputs.node.names.copy(),
+                defn=new_outputs_defn,
+                module_name=node_info.module_name,
+            )
+
+            base_result_sym = base_node_outputs.node.names[attribute_name].node
+            if isinstance(base_result_sym, Var):
+                new_result_sym = Var.deserialize(base_result_sym.serialize())
+                new_result_sym._fullname = f"{new_outputs_fullname}.{attribute_name}"
+                new_result_sym.type = base_node_resolved_type
+                new_outputs_sym.names[attribute_name].node = new_result_sym
+                new_outputs_sym.bases.append(Instance(base_node_outputs.node, []))
+                node_info.names["Outputs"].node = new_outputs_sym
 
     def _base_node_class_hook(self, ctx: ClassDefContext) -> None:
         """
@@ -472,10 +487,22 @@ class VellumMypyPlugin(Plugin):
 
     def _run_method_hook(self, ctx: MethodContext) -> MypyType:
         """
-        We use this to target `Workflow.run()` so that the WorkflowExecutionFulfilledEvent is properly typed
-        using the `Outputs` class defined on the user-defined subclass of `Workflow`.
+        We use this to target:
+        - `BaseWorkflow.run()`, so that the WorkflowExecutionFulfilledEvent is properly typed
+           using the `Outputs` class defined on the user-defined subclass of `Workflow`.
+        - `BaseNode.run()`, so that the `Outputs` class defined on the user-defined subclass of `Node`
+          is properly typed.
         """
 
+        if isinstance(ctx.default_return_type, TypeAliasType):
+            return self._workflow_run_method_hook(ctx)
+
+        if isinstance(ctx.default_return_type, Instance):
+            return self._node_run_method_hook(ctx)
+
+        return ctx.default_return_type
+
+    def _workflow_run_method_hook(self, ctx: MethodContext) -> MypyType:
         if not isinstance(ctx.default_return_type, TypeAliasType):
             return ctx.default_return_type
 
@@ -494,7 +521,7 @@ class VellumMypyPlugin(Plugin):
         if fulfilled_event.type.fullname != "vellum.workflows.events.workflow.WorkflowExecutionFulfilledEvent":
             return ctx.default_return_type
 
-        outputs_node = self._get_outputs_node(ctx)
+        outputs_node = self._get_workflow_outputs_type_info(ctx)
         if not outputs_node:
             return ctx.default_return_type
 
@@ -513,6 +540,19 @@ class VellumMypyPlugin(Plugin):
             column=ctx.default_return_type.column,
         )
 
+    def _node_run_method_hook(self, ctx: MethodContext) -> MypyType:
+        if not isinstance(ctx.default_return_type, Instance):
+            return ctx.default_return_type
+
+        if not _is_subclass(ctx.default_return_type.type, "vellum.workflows.nodes.bases.base.BaseNode.Outputs"):
+            return ctx.default_return_type
+
+        outputs_node = self._get_node_outputs_type_info(ctx)
+        if not outputs_node:
+            return ctx.default_return_type
+
+        return Instance(outputs_node, [])
+
     def _stream_method_hook(self, ctx: MethodContext) -> MypyType:
         """
         We use this to target `Workflow.stream()` so that the WorkflowExecutionFulfilledEvent is properly typed
@@ -557,7 +597,7 @@ class VellumMypyPlugin(Plugin):
         if fulfilled_event_index == -1 or not fulfilled_event:
             return ctx.default_return_type
 
-        outputs_node = self._get_outputs_node(ctx)
+        outputs_node = self._get_workflow_outputs_type_info(ctx)
         if not outputs_node:
             return ctx.default_return_type
 
@@ -594,7 +634,7 @@ class VellumMypyPlugin(Plugin):
             column=ctx.default_return_type.column,
         )
 
-    def _get_outputs_node(self, ctx: MethodContext) -> Optional[TypeInfo]:
+    def _get_workflow_outputs_type_info(self, ctx: MethodContext) -> Optional[TypeInfo]:
         if not isinstance(ctx.context, CallExpr):
             return None
 
@@ -624,6 +664,35 @@ class VellumMypyPlugin(Plugin):
 
         return resolved_outputs_node.node
 
+    def _get_node_outputs_type_info(self, ctx: MethodContext) -> Optional[TypeInfo]:
+        if not isinstance(ctx.context, CallExpr):
+            return None
+
+        if not isinstance(ctx.context.callee, MemberExpr):
+            return None
+
+        expr = ctx.context.callee.expr
+        instance = ctx.api.get_expression_type(expr)
+        if not isinstance(instance, Instance) or not _is_subclass(
+            instance.type, "vellum.workflows.nodes.bases.base.BaseNode"
+        ):
+            return None
+
+        outputs_node = instance.type.names.get("Outputs")
+
+        if (
+            not outputs_node
+            or not isinstance(outputs_node.node, TypeInfo)
+            or not _is_subclass(outputs_node.node, "vellum.workflows.outputs.base.BaseOutputs")
+        ):
+            return None
+
+        # TODO: For some reason, returning the correct Outputs type info is causing `result` to not
+        # be found. `test_templating_node.py` is the best place to test this.
+        # https://app.shortcut.com/vellum/story/6132
+        # return outputs_node.node
+        return None
+
     def _resolve_descriptors_in_outputs(self, type_info: SymbolTableNode) -> SymbolTableNode:
         new_type_info = type_info.copy()
         if not isinstance(new_type_info.node, TypeInfo):

vellum/utils/templating/render.py

@@ -28,6 +28,9 @@ def render_sandboxed_jinja_template(
             keep_trailing_newline=True,
             finalize=finalize,
         )
+        environment.policies["json.dumps_kwargs"] = {
+            "cls": DefaultStateEncoder,
+        }
 
         if jinja_custom_filters:
             environment.filters.update(jinja_custom_filters)

vellum/workflows/nodes/bases/base.py

@@ -324,6 +324,10 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
             if not descriptor.instance:
                 continue
 
+            if any(isinstance(t, type) and issubclass(t, BaseDescriptor) for t in descriptor.types):
+                # We don't want to resolve attributes that are _meant_ to be descriptors
+                continue
+
             resolved_value = resolve_value(descriptor.instance, self.state, path=descriptor.name, memo=inputs)
             setattr(self, descriptor.name, resolved_value)
 

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

@@ -1,5 +1,7 @@
 from typing import Callable, Generic, Optional, Type
 
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.utils import resolve_value
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
@@ -21,6 +23,7 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
 
     max_attempts: int
     retry_on_error_code: Optional[WorkflowErrorCode] = None
+    retry_on_condition: Optional[BaseDescriptor] = None
 
     class SubworkflowInputs(BaseInputs):
         attempt_number: int
@@ -55,18 +58,36 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
                 last_exception = NodeException(
                     code=WorkflowErrorCode.INVALID_OUTPUTS,
                     message=f"""Unexpected rejection on attempt {attempt_number}: {terminal_event.error.code.value}.
+Message: {terminal_event.error.message}""",
+                )
+                break
+            elif self.retry_on_condition and not resolve_value(self.retry_on_condition, self.state):
+                last_exception = NodeException(
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
+                    message=f"""Rejection failed on attempt {attempt_number}: {terminal_event.error.code.value}.
 Message: {terminal_event.error.message}""",
                 )
                 break
             else:
-                last_exception = Exception(terminal_event.error.message)
+                last_exception = NodeException(
+                    terminal_event.error.message,
+                    code=terminal_event.error.code,
+                )
 
         raise last_exception
 
     @classmethod
     def wrap(
-        cls, max_attempts: int, retry_on_error_code: Optional[WorkflowErrorCode] = None
+        cls,
+        max_attempts: int,
+        retry_on_error_code: Optional[WorkflowErrorCode] = None,
+        retry_on_condition: Optional[BaseDescriptor] = None,
     ) -> Callable[..., Type["RetryNode"]]:
         return create_adornment(
-            cls, attributes={"max_attempts": max_attempts, "retry_on_error_code": retry_on_error_code}
+            cls,
+            attributes={
+                "max_attempts": max_attempts,
+                "retry_on_error_code": retry_on_error_code,
+                "retry_on_condition": retry_on_condition,
+            },
         )

vellum/workflows/nodes/core/retry_node/tests/test_node.py

@@ -6,6 +6,7 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state.base import BaseState, StateMeta
 
 
@@ -91,3 +92,42 @@ def test_retry_node__use_parent_inputs_and_state():
 
     # THEN the data is used successfully
     assert outputs.value == "foo bar"
+
+
+def test_retry_node__condition_arg_successfully_retries():
+    # GIVEN workflow Inputs and State
+    class State(BaseState):
+        count = 0
+
+    # AND a retry node that retries on a condition
+    @RetryNode.wrap(
+        max_attempts=5,
+        retry_on_condition=LazyReference(lambda: State.count.less_than(3)),
+    )
+    class TestNode(BaseNode[State]):
+        attempt_number = RetryNode.SubworkflowInputs.attempt_number
+
+        class Outputs(BaseOutputs):
+            value: str
+
+        def run(self) -> Outputs:
+            if not isinstance(self.state.meta.parent, State):
+                raise NodeException(message="Failed to resolve parent state")
+
+            self.state.meta.parent.count += 1
+            raise NodeException(message=f"This is failure attempt {self.attempt_number}")
+
+    # WHEN the node is run
+    node = TestNode(state=State())
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the exception raised is the last one
+    assert (
+        exc_info.value.message
+        == """Rejection failed on attempt 3: INTERNAL_ERROR.
+Message: This is failure attempt 3"""
+    )
+
+    # AND the state was updated each time
+    assert node.state.count == 3

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

@@ -1,5 +1,6 @@
 import json
 
+from vellum.client.types.function_call import FunctionCall
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.state import BaseState
@@ -21,7 +22,9 @@ def test_templating_node__dict_output():
     outputs = node.run()
 
     # THEN the output is json serializable
-    assert json.loads(outputs.result) == {"key": "value"}
+    # https://app.shortcut.com/vellum/story/6132
+    dump: str = outputs.result  # type: ignore[assignment]
+    assert json.loads(dump) == {"key": "value"}
 
 
 def test_templating_node__int_output():
@@ -106,3 +109,19 @@ def test_templating_node__execution_count_reference():
 
     # THEN the output is just the total
     assert outputs.result == "0"
+
+
+def test_templating_node__pydantic_to_json():
+    # GIVEN a templating node that uses tojson on a pydantic model
+    class JSONTemplateNode(TemplatingNode[BaseState, Json]):
+        template = "{{ function_call | tojson }}"
+        inputs = {
+            "function_call": FunctionCall(name="test", arguments={"key": "value"}),
+        }
+
+    # WHEN the node is run
+    node = JSONTemplateNode()
+    outputs = node.run()
+
+    # THEN the output is the expected JSON
+    assert outputs.result == {"name": "test", "arguments": {"key": "value"}, "id": None}

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

@@ -2,7 +2,7 @@ from typing import Optional, Union
 
 from vellum.workflows.constants import AuthorizationType
 from vellum.workflows.nodes.displayable.bases.api_node import BaseAPINode
-from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.types.core import VellumSecret
 
 
 class APINode(BaseAPINode):
@@ -24,8 +24,8 @@ class APINode(BaseAPINode):
 
     authorization_type: Optional[AuthorizationType] = None
     api_key_header_key: Optional[str] = None
-    api_key_header_value: Optional[Union[str, VellumSecretReference]] = None
-    bearer_token_value: Optional[Union[str, VellumSecretReference]] = None
+    api_key_header_value: Optional[Union[str, VellumSecret]] = None
+    bearer_token_value: Optional[Union[str, VellumSecret]] = None
 
     def run(self) -> BaseAPINode.Outputs:
         headers = self.headers or {}

{vellum_ai-0.12.16.dist-info → vellum_ai-0.12.17.dist-info}/METADATA

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

{vellum_ai-0.12.16.dist-info → vellum_ai-0.12.17.dist-info}/RECORD

@@ -5,14 +5,14 @@ vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3
 vellum_cli/config.py,sha256=998IZbvkrw2avjbvs8bw6NrbEgGz5UBKRbvKAcastJg,5493
 vellum_cli/image_push.py,sha256=SJwhwWJsLjwGNezNVd_oCVpFMfPsAB3dfLWmriZZUtw,4419
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
-vellum_cli/pull.py,sha256=hYQBe2_-Y4Ieh_Jo54EwBUpqPYq6UYitUBPgqyOVba0,7499
+vellum_cli/pull.py,sha256=zf0y22XptUYI_hMP_4Q1CEo9s2wALsTJcCXNd-_ibd8,7551
 vellum_cli/push.py,sha256=iQ2H-vY8njqiYgIifGooqopqkb1BUUA3I7IN7VIHKv8,6149
 vellum_cli/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_cli/tests/conftest.py,sha256=eFGwBxib3Nki830lIFintB0b6r4x8T_KMnmzhlTY5x0,1337
+vellum_cli/tests/conftest.py,sha256=Jv-QUjXoCDhsvVvXEjOenNqRArO_xXhtNuCYt4wiOyE,1421
 vellum_cli/tests/test_config.py,sha256=uvKGDc8BoVyT9_H0Z-g8469zVxomn6Oi3Zj-vK7O_wU,2631
 vellum_cli/tests/test_main.py,sha256=qDZG-aQauPwBwM6A2DIu1494n47v3pL28XakTbLGZ-k,272
 vellum_cli/tests/test_pull.py,sha256=6gbASF6ASy5YcdWjOCt6b5K0u2VWsFegdrTWu6sEVKs,19613
-vellum_cli/tests/test_push.py,sha256=sqPy1N9TOa0Wwk4tfEZm18BnfcwRSoQboTWW_PHQA-E,8145
+vellum_cli/tests/test_push.py,sha256=8o8DFW9HCakhfZMS1Iql13RC90hF9pM630Y-rQbo234,8593
 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
@@ -84,7 +84,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=7yb5YnhvHQUJusa1WyUZcAKGol3z-Lfu07EfD03eEW4,111291
 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=2N1VoBiQAederYuBKuAwXoDAkbTGPnkciOAPvOv1Xyk,1869
+vellum/client/core/client_wrapper.py,sha256=Ip_JnKYTedcuVV4puh0iLeYLzPtHAxgnmwqZyCKXSCY,1869
 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
@@ -668,7 +668,7 @@ vellum/evaluations/utils/paginator.py,sha256=rEED_BJAXAM6tM1yMwHePNzszjq_tTq4NbQ
 vellum/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/plugins/pydantic.py,sha256=dNtZWHo-IdseG52C2RoTanxyTJg0AhPZrH-9lbNqwYg,2604
 vellum/plugins/utils.py,sha256=U9ZY9KdE3RRvbcG01hXxu9CvfJD6Fo7nJDgcHjQn0FI,606
-vellum/plugins/vellum_mypy.py,sha256=VC15EzjTsXOb9uF1bky4rcxePP-0epMVmCsLB2z4Dh8,24816
+vellum/plugins/vellum_mypy.py,sha256=7JmyuTX723-XWnWoE6afiUNOkHyAufCUZwxck9BP2aI,27784
 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
@@ -1221,7 +1221,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=0vgkwhu2A6o64aT4fUdTSLFCEMbeRjAKAuvv2k2LYGY,1772
+vellum/utils/templating/render.py,sha256=OwdZTJlQv_qsygMX3LOFr4d1_2oz3r9vGKyj7cX05VE,1876
 vellum/utils/typing.py,sha256=wx_daFqD69cYkuJTVnvNrpjhqC3uuhbnyJ9_bIwC9OU,327
 vellum/utils/uuid.py,sha256=Ch6wWRgwICxLxJCTl5iE3EdRlZj2zADR-zUMUtjcMWM,214
 vellum/version.py,sha256=jq-1PlAYxN9AXuaZqbYk9ak27SgE2lw9Ia5gx1b1gVI,76
@@ -1286,7 +1286,7 @@ vellum/workflows/inputs/base.py,sha256=1kMgr0WqCYdWUqgFvgSoAMw2067FAlgwhGXLgbIOr
 vellum/workflows/logging.py,sha256=_a217XogktV4Ncz6xKFz7WfYmZAzkfVRVuC0rWob8ls,437
 vellum/workflows/nodes/__init__.py,sha256=aVdQVv7Y3Ro3JlqXGpxwaU2zrI06plDHD2aumH5WUIs,1157
 vellum/workflows/nodes/bases/__init__.py,sha256=cniHuz_RXdJ4TQgD8CBzoiKDiPxg62ErdVpCbWICX64,58
-vellum/workflows/nodes/bases/base.py,sha256=RREFzYPxemKUvQc0NfnwQmby-p_BE3O-TbVWKbQFdfs,14271
+vellum/workflows/nodes/bases/base.py,sha256=CcWQBMR3cx5vftqQp_oJ_GncULJIbOLMyNP4KZNgU10,14487
 vellum/workflows/nodes/bases/base_adornment_node.py,sha256=eFTgsPCYb3eyGS0-kw7C6crFnwFx437R5wh9-8bWYts,2905
 vellum/workflows/nodes/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/bases/tests/test_base_node.py,sha256=51CueFVty9XYASC0rKr1cXWejho5WElmhfhp6cCONy0,3811
@@ -1302,19 +1302,19 @@ vellum/workflows/nodes/core/map_node/node.py,sha256=DTMoGqtR8MyfZ8jy8apNoN-4KFFF
 vellum/workflows/nodes/core/map_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/map_node/tests/test_node.py,sha256=RHSZs7t6mW3UWvRrXnHZqaXVdRT2ZquOK_YHJ-gzXsU,1871
 vellum/workflows/nodes/core/retry_node/__init__.py,sha256=lN2bIy5a3Uzhs_FYCrooADyYU6ZGShtvLKFWpelwPvo,60
-vellum/workflows/nodes/core/retry_node/node.py,sha256=lAABgo2E_pWkzOYUBGzC1SnywgwtGuJojwT602fKCUc,3153
+vellum/workflows/nodes/core/retry_node/node.py,sha256=QEpxhKOyxDkRoAn2b0PToZWtAGQetSQYVTpb9yCOLlw,4028
 vellum/workflows/nodes/core/retry_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/nodes/core/retry_node/tests/test_node.py,sha256=QXTnHwmJHISxXjvZMeuuEo0iVugVMJyaJoggI8yKXfI,3132
+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=N-NOBd-UY91qO9orCcW4KEbhNvDQivZPA-PCxs-M0RM,4204
-vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=nW_kyJ9RAqz45_kJE_rlhOOvbV4OO3hecP-P-ydQpkw,2845
+vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=3rN5Ncwq87Y2YKyobNXALCih6OnUTJf55otnR4V20C8,3557
 vellum/workflows/nodes/core/try_node/__init__.py,sha256=JVD4DrldTIqFQQFrubs9KtWCCc0YCAc7Fzol5ZWIWeM,56
 vellum/workflows/nodes/core/try_node/node.py,sha256=_lTmSYCiz7lktaxpNWUCglNi8_5Sfy8Rpiov5SeKVMw,3920
 vellum/workflows/nodes/core/try_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/try_node/tests/test_node.py,sha256=Wc2kLl-MkffsBxl3IiFaqLd16e2Iosxhk7qBnojPvQg,4092
 vellum/workflows/nodes/displayable/__init__.py,sha256=6F_4DlSwvHuilWnIalp8iDjjDXl0Nmz4QzJV2PYe5RI,1023
 vellum/workflows/nodes/displayable/api_node/__init__.py,sha256=MoxdQSnidIj1Nf_d-hTxlOxcZXaZnsWFDbE-PkTK24o,56
-vellum/workflows/nodes/displayable/api_node/node.py,sha256=ehvJNkI-L_WLn8pszCaRkQuhRDgioCLaX6TT72KM_9E,2111
+vellum/workflows/nodes/displayable/api_node/node.py,sha256=ocXjg0Sp52V6DllTDZXUgOp9qBw6l_KbtWcaxowsDM4,2070
 vellum/workflows/nodes/displayable/bases/__init__.py,sha256=0mWIx3qUrzllV7jqt7wN03vWGMuI1WrrLZeMLT2Cl2c,304
 vellum/workflows/nodes/displayable/bases/api_node/__init__.py,sha256=1jwx4WC358CLA1jgzl_UD-rZmdMm2v9Mps39ndwCD7U,64
 vellum/workflows/nodes/displayable/bases/api_node/node.py,sha256=ywPwCxDesgnE86Wjyf2ZyG2Dr2O7xW8D4tPHZB5Se_s,2477
@@ -1415,8 +1415,8 @@ vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
 vellum/workflows/workflows/base.py,sha256=qdZYQq-jjdr0fYT0FCfmFuI5ypE3pANupgYcOqqML0o,18884
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
-vellum_ai-0.12.16.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.12.16.dist-info/METADATA,sha256=dSVFM4nOL_e1hiNJX-uVrRBdO4-1wDtqhA3MZKwMLTM,5328
-vellum_ai-0.12.16.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.12.16.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.12.16.dist-info/RECORD,,
+vellum_ai-0.12.17.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.12.17.dist-info/METADATA,sha256=01ujt39Or2zqLaKV_VsarQ9H_hkgSc1YfrsEOMokUUk,5328
+vellum_ai-0.12.17.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.12.17.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.12.17.dist-info/RECORD,,

vellum_cli/pull.py

@@ -173,10 +173,13 @@ def pull_command(
                     target.write(content)
 
     if metadata_json:
-        workflow_config.container_image_name = metadata_json.get("runner_config", {}).get("container_image_name")
-        workflow_config.container_image_tag = metadata_json.get("runner_config", {}).get("container_image_tag")
-        if workflow_config.container_image_name and not workflow_config.container_image_tag:
-            workflow_config.container_image_tag = "latest"
+        runner_config = metadata_json.get("runner_config")
+
+        if runner_config:
+            workflow_config.container_image_name = runner_config.get("container_image_name")
+            workflow_config.container_image_tag = runner_config.get("container_image_tag")
+            if workflow_config.container_image_name and not workflow_config.container_image_tag:
+                workflow_config.container_image_tag = "latest"
 
     if include_json:
         logger.warning(

vellum_cli/tests/conftest.py

@@ -18,11 +18,13 @@ class MockModuleResult:
 
 
 @pytest.fixture
-def mock_module() -> Generator[MockModuleResult, None, None]:
+def mock_module(request) -> Generator[MockModuleResult, None, None]:
     current_dir = os.getcwd()
     temp_dir = tempfile.mkdtemp()
     os.chdir(temp_dir)
-    module = "examples.mock"
+
+    # Use the test name to create a unique module path
+    module = f"examples.mock.{request.node.name}"
     workflow_sandbox_id = str(uuid4())
 
     def set_pyproject_toml(vellum_config: Dict[str, Any]) -> None:

vellum_cli/tests/test_push.py

@@ -110,13 +110,17 @@ class ExampleWorkflow(BaseWorkflow):
     # THEN it should succeed
     assert result.exit_code == 0
 
+    # Get the last part of the module path and format it
+    expected_label = mock_module.module.split(".")[-1].replace("_", " ").title()
+    expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
+
     # AND we should have called the push API with the correct args
     vellum_client.workflows.push.assert_called_once()
     call_args = vellum_client.workflows.push.call_args.kwargs
     assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
-    assert call_args["label"] == "Mock"
+    assert call_args["label"] == expected_label
     assert is_valid_uuid(call_args["workflow_sandbox_id"])
-    assert call_args["artifact"].name == "examples__mock.tar.gz"
+    assert call_args["artifact"].name == expected_artifact_name
     assert "deplyment_config" not in call_args
 
     extracted_files = _extract_tar_gz(call_args["artifact"].read())
@@ -160,13 +164,17 @@ class ExampleWorkflow(BaseWorkflow):
     # THEN it should succeed
     assert result.exit_code == 0
 
+    # Get the last part of the module path and format it
+    expected_label = mock_module.module.split(".")[-1].replace("_", " ").title()
+    expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
+
     # AND we should have called the push API with the correct args
     vellum_client.workflows.push.assert_called_once()
     call_args = vellum_client.workflows.push.call_args.kwargs
     assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
-    assert call_args["label"] == "Mock"
+    assert call_args["label"] == expected_label
     assert is_valid_uuid(call_args["workflow_sandbox_id"])
-    assert call_args["artifact"].name == "examples__mock.tar.gz"
+    assert call_args["artifact"].name == expected_artifact_name
     assert call_args["deployment_config"] == "{}"
 
     extracted_files = _extract_tar_gz(call_args["artifact"].read())
@@ -231,6 +239,6 @@ class ExampleWorkflow(BaseWorkflow):
 
     # AND the report should be in the output
     assert "## Errors" in result.output
-    # assert "Serialization is not supported." in result.output
+    assert "Serialization is not supported." in result.output
     assert "## Proposed Diffs" in result.output
     assert "iterable_item_added" in result.output