vellum-ai 0.12.16__py3-none-any.whl → 0.13.0__py3-none-any.whl

vellum/client/types/__init__.py

@@ -14,14 +14,11 @@ from .array_chat_message_content_item import ArrayChatMessageContentItem
 from .array_chat_message_content_item_request import ArrayChatMessageContentItemRequest
 from .array_chat_message_content_request import ArrayChatMessageContentRequest
 from .array_input import ArrayInput
-from .array_variable_value import ArrayVariableValue
-from .array_variable_value_item import ArrayVariableValueItem
 from .array_vellum_value import ArrayVellumValue
 from .array_vellum_value_request import ArrayVellumValueRequest
 from .audio_chat_message_content import AudioChatMessageContent
 from .audio_chat_message_content_request import AudioChatMessageContentRequest
 from .audio_prompt_block import AudioPromptBlock
-from .audio_variable_value import AudioVariableValue
 from .audio_vellum_value import AudioVellumValue
 from .audio_vellum_value_request import AudioVellumValueRequest
 from .basic_vectorizer_intfloat_multilingual_e_5_large import BasicVectorizerIntfloatMultilingualE5Large
@@ -40,7 +37,6 @@ from .basic_vectorizer_sentence_transformers_multi_qa_mpnet_base_dot_v_1_request
 )
 from .chat_history_input import ChatHistoryInput
 from .chat_history_input_request import ChatHistoryInputRequest
-from .chat_history_variable_value import ChatHistoryVariableValue
 from .chat_history_vellum_value import ChatHistoryVellumValue
 from .chat_history_vellum_value_request import ChatHistoryVellumValueRequest
 from .chat_message import ChatMessage
@@ -97,7 +93,6 @@ from .environment_enum import EnvironmentEnum
 from .ephemeral_prompt_cache_config import EphemeralPromptCacheConfig
 from .ephemeral_prompt_cache_config_type_enum import EphemeralPromptCacheConfigTypeEnum
 from .error_input import ErrorInput
-from .error_variable_value import ErrorVariableValue
 from .error_vellum_value import ErrorVellumValue
 from .error_vellum_value_request import ErrorVellumValueRequest
 from .execute_prompt_event import ExecutePromptEvent
@@ -141,7 +136,6 @@ from .function_call_chat_message_content_value_request import FunctionCallChatMe
 from .function_call_input import FunctionCallInput
 from .function_call_prompt_block import FunctionCallPromptBlock
 from .function_call_request import FunctionCallRequest
-from .function_call_variable_value import FunctionCallVariableValue
 from .function_call_vellum_value import FunctionCallVellumValue
 from .function_call_vellum_value_request import FunctionCallVellumValueRequest
 from .function_definition import FunctionDefinition
@@ -169,7 +163,6 @@ from .hkunlp_instructor_xl_vectorizer_request import HkunlpInstructorXlVectorize
 from .image_chat_message_content import ImageChatMessageContent
 from .image_chat_message_content_request import ImageChatMessageContentRequest
 from .image_prompt_block import ImagePromptBlock
-from .image_variable_value import ImageVariableValue
 from .image_vellum_value import ImageVellumValue
 from .image_vellum_value_request import ImageVellumValueRequest
 from .indexing_config_vectorizer import IndexingConfigVectorizer
@@ -185,7 +178,6 @@ from .iteration_state_enum import IterationStateEnum
 from .jinja_prompt_block import JinjaPromptBlock
 from .json_input import JsonInput
 from .json_input_request import JsonInputRequest
-from .json_variable_value import JsonVariableValue
 from .json_vellum_value import JsonVellumValue
 from .json_vellum_value_request import JsonVellumValueRequest
 from .logical_operator import LogicalOperator
@@ -226,6 +218,7 @@ from .named_test_case_string_variable_value import NamedTestCaseStringVariableVa
 from .named_test_case_string_variable_value_request import NamedTestCaseStringVariableValueRequest
 from .named_test_case_variable_value import NamedTestCaseVariableValue
 from .named_test_case_variable_value_request import NamedTestCaseVariableValueRequest
+from .new_member_join_behavior_enum import NewMemberJoinBehaviorEnum
 from .node_input_compiled_array_value import NodeInputCompiledArrayValue
 from .node_input_compiled_chat_history_value import NodeInputCompiledChatHistoryValue
 from .node_input_compiled_error_value import NodeInputCompiledErrorValue
@@ -247,7 +240,6 @@ from .node_output_compiled_value import NodeOutputCompiledValue
 from .normalized_log_probs import NormalizedLogProbs
 from .normalized_token_log_probs import NormalizedTokenLogProbs
 from .number_input import NumberInput
-from .number_variable_value import NumberVariableValue
 from .number_vellum_value import NumberVellumValue
 from .number_vellum_value_request import NumberVellumValueRequest
 from .open_ai_vectorizer_config import OpenAiVectorizerConfig
@@ -258,6 +250,7 @@ from .open_ai_vectorizer_text_embedding_3_small import OpenAiVectorizerTextEmbed
 from .open_ai_vectorizer_text_embedding_3_small_request import OpenAiVectorizerTextEmbedding3SmallRequest
 from .open_ai_vectorizer_text_embedding_ada_002 import OpenAiVectorizerTextEmbeddingAda002
 from .open_ai_vectorizer_text_embedding_ada_002_request import OpenAiVectorizerTextEmbeddingAda002Request
+from .organization_read import OrganizationRead
 from .paginated_container_image_read_list import PaginatedContainerImageReadList
 from .paginated_deployment_release_tag_read_list import PaginatedDeploymentReleaseTagReadList
 from .paginated_document_index_read_list import PaginatedDocumentIndexReadList
@@ -321,7 +314,6 @@ from .search_result_meta import SearchResultMeta
 from .search_result_meta_request import SearchResultMetaRequest
 from .search_result_request import SearchResultRequest
 from .search_results_input import SearchResultsInput
-from .search_results_variable_value import SearchResultsVariableValue
 from .search_results_vellum_value import SearchResultsVellumValue
 from .search_results_vellum_value_request import SearchResultsVellumValueRequest
 from .search_weights_request import SearchWeightsRequest
@@ -342,7 +334,6 @@ from .string_chat_message_content import StringChatMessageContent
 from .string_chat_message_content_request import StringChatMessageContentRequest
 from .string_input import StringInput
 from .string_input_request import StringInputRequest
-from .string_variable_value import StringVariableValue
 from .string_vellum_value import StringVellumValue
 from .string_vellum_value_request import StringVellumValueRequest
 from .submit_completion_actual_request import SubmitCompletionActualRequest
@@ -522,6 +513,7 @@ from .workflow_result_event_output_data_number import WorkflowResultEventOutputD
 from .workflow_result_event_output_data_search_results import WorkflowResultEventOutputDataSearchResults
 from .workflow_result_event_output_data_string import WorkflowResultEventOutputDataString
 from .workflow_stream_event import WorkflowStreamEvent
+from .workspace_read import WorkspaceRead
 from .workspace_secret_read import WorkspaceSecretRead
 
 __all__ = [
@@ -539,14 +531,11 @@ __all__ = [
     "ArrayChatMessageContentItemRequest",
     "ArrayChatMessageContentRequest",
     "ArrayInput",
-    "ArrayVariableValue",
-    "ArrayVariableValueItem",
     "ArrayVellumValue",
     "ArrayVellumValueRequest",
     "AudioChatMessageContent",
     "AudioChatMessageContentRequest",
     "AudioPromptBlock",
-    "AudioVariableValue",
     "AudioVellumValue",
     "AudioVellumValueRequest",
     "BasicVectorizerIntfloatMultilingualE5Large",
@@ -557,7 +546,6 @@ __all__ = [
     "BasicVectorizerSentenceTransformersMultiQaMpnetBaseDotV1Request",
     "ChatHistoryInput",
     "ChatHistoryInputRequest",
-    "ChatHistoryVariableValue",
     "ChatHistoryVellumValue",
     "ChatHistoryVellumValueRequest",
     "ChatMessage",
@@ -614,7 +602,6 @@ __all__ = [
     "EphemeralPromptCacheConfig",
     "EphemeralPromptCacheConfigTypeEnum",
     "ErrorInput",
-    "ErrorVariableValue",
     "ErrorVellumValue",
     "ErrorVellumValueRequest",
     "ExecutePromptEvent",
@@ -658,7 +645,6 @@ __all__ = [
     "FunctionCallInput",
     "FunctionCallPromptBlock",
     "FunctionCallRequest",
-    "FunctionCallVariableValue",
     "FunctionCallVellumValue",
     "FunctionCallVellumValueRequest",
     "FunctionDefinition",
@@ -682,7 +668,6 @@ __all__ = [
     "ImageChatMessageContent",
     "ImageChatMessageContentRequest",
     "ImagePromptBlock",
-    "ImageVariableValue",
     "ImageVellumValue",
     "ImageVellumValueRequest",
     "IndexingConfigVectorizer",
@@ -698,7 +683,6 @@ __all__ = [
     "JinjaPromptBlock",
     "JsonInput",
     "JsonInputRequest",
-    "JsonVariableValue",
     "JsonVellumValue",
     "JsonVellumValueRequest",
     "LogicalOperator",
@@ -739,6 +723,7 @@ __all__ = [
     "NamedTestCaseStringVariableValueRequest",
     "NamedTestCaseVariableValue",
     "NamedTestCaseVariableValueRequest",
+    "NewMemberJoinBehaviorEnum",
     "NodeInputCompiledArrayValue",
     "NodeInputCompiledChatHistoryValue",
     "NodeInputCompiledErrorValue",
@@ -760,7 +745,6 @@ __all__ = [
     "NormalizedLogProbs",
     "NormalizedTokenLogProbs",
     "NumberInput",
-    "NumberVariableValue",
     "NumberVellumValue",
     "NumberVellumValueRequest",
     "OpenAiVectorizerConfig",
@@ -771,6 +755,7 @@ __all__ = [
     "OpenAiVectorizerTextEmbedding3SmallRequest",
     "OpenAiVectorizerTextEmbeddingAda002",
     "OpenAiVectorizerTextEmbeddingAda002Request",
+    "OrganizationRead",
     "PaginatedContainerImageReadList",
     "PaginatedDeploymentReleaseTagReadList",
     "PaginatedDocumentIndexReadList",
@@ -834,7 +819,6 @@ __all__ = [
     "SearchResultMetaRequest",
     "SearchResultRequest",
     "SearchResultsInput",
-    "SearchResultsVariableValue",
     "SearchResultsVellumValue",
     "SearchResultsVellumValueRequest",
     "SearchWeightsRequest",
@@ -855,7 +839,6 @@ __all__ = [
     "StringChatMessageContentRequest",
     "StringInput",
     "StringInputRequest",
-    "StringVariableValue",
     "StringVellumValue",
     "StringVellumValueRequest",
     "SubmitCompletionActualRequest",
@@ -1019,5 +1002,6 @@ __all__ = [
     "WorkflowResultEventOutputDataSearchResults",
     "WorkflowResultEventOutputDataString",
     "WorkflowStreamEvent",
+    "WorkspaceRead",
     "WorkspaceSecretRead",
 ]

vellum/client/types/new_member_join_behavior_enum.py

@@ -0,0 +1,8 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+NewMemberJoinBehaviorEnum = typing.Union[
+    typing.Literal["AUTO_ACCEPT_FROM_SHARED_DOMAIN", "ALLOW_REQUESTS_FROM_SHARED_DOMAIN", "REQUIRE_EXPLICIT_INVITE"],
+    typing.Any,
+]

vellum/client/types/{function_call_variable_value.py → organization_read.py}

@@ -2,14 +2,16 @@
 
 from ..core.pydantic_utilities import UniversalBaseModel
 import typing
-from .function_call import FunctionCall
+from .new_member_join_behavior_enum import NewMemberJoinBehaviorEnum
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
 
 
-class FunctionCallVariableValue(UniversalBaseModel):
-    type: typing.Literal["FUNCTION_CALL"] = "FUNCTION_CALL"
-    value: typing.Optional[FunctionCall] = None
+class OrganizationRead(UniversalBaseModel):
+    id: str
+    name: str
+    allow_staff_access: typing.Optional[bool] = None
+    new_member_join_behavior: NewMemberJoinBehaviorEnum
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

vellum/client/types/workflow_execution_workflow_result_event.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
-from .array_variable_value import ArrayVariableValue
 from .array_vellum_value import ArrayVellumValue
 import typing
 from .workflow_result_event import WorkflowResultEvent
@@ -32,5 +31,4 @@ class WorkflowExecutionWorkflowResultEvent(UniversalBaseModel):
             extra = pydantic.Extra.allow
 
 
-update_forward_refs(ArrayVariableValue, WorkflowExecutionWorkflowResultEvent=WorkflowExecutionWorkflowResultEvent)
 update_forward_refs(ArrayVellumValue, WorkflowExecutionWorkflowResultEvent=WorkflowExecutionWorkflowResultEvent)

vellum/client/types/workflow_result_event.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
-from .array_variable_value import ArrayVariableValue
 from .array_vellum_value import ArrayVellumValue
 from .workflow_node_result_event_state import WorkflowNodeResultEventState
 import datetime as dt
@@ -35,5 +34,4 @@ class WorkflowResultEvent(UniversalBaseModel):
             extra = pydantic.Extra.allow
 
 
-update_forward_refs(ArrayVariableValue, WorkflowResultEvent=WorkflowResultEvent)
 update_forward_refs(ArrayVellumValue, WorkflowResultEvent=WorkflowResultEvent)

vellum/client/types/workflow_result_event_output_data_array.py

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
-from .array_variable_value import ArrayVariableValue
+from .array_vellum_value import ArrayVellumValue
 import typing
 from .workflow_node_result_event_state import WorkflowNodeResultEventState
 import pydantic
-from .array_variable_value_item import ArrayVariableValueItem
+from .vellum_value import VellumValue
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 from ..core.pydantic_utilities import update_forward_refs
 
@@ -26,7 +26,7 @@ class WorkflowResultEventOutputDataArray(UniversalBaseModel):
     """
 
     type: typing.Literal["ARRAY"] = "ARRAY"
-    value: typing.Optional[typing.List[ArrayVariableValueItem]] = None
+    value: typing.Optional[typing.List[VellumValue]] = None
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
@@ -38,4 +38,4 @@ class WorkflowResultEventOutputDataArray(UniversalBaseModel):
             extra = pydantic.Extra.allow
 
 
-update_forward_refs(ArrayVariableValue, WorkflowResultEventOutputDataArray=WorkflowResultEventOutputDataArray)
+update_forward_refs(ArrayVellumValue, WorkflowResultEventOutputDataArray=WorkflowResultEventOutputDataArray)

vellum/client/types/{string_variable_value.py → workspace_read.py}

@@ -1,14 +1,21 @@
 # This file was auto-generated by Fern from our API Definition.
 
 from ..core.pydantic_utilities import UniversalBaseModel
-import typing
-from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
+import datetime as dt
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+
 
+class WorkspaceRead(UniversalBaseModel):
+    id: str
+    name: str = pydantic.Field()
+    """
+    The name of the Workspace.
+    """
 
-class StringVariableValue(UniversalBaseModel):
-    type: typing.Literal["STRING"] = "STRING"
-    value: typing.Optional[str] = None
+    label: str
+    created: dt.datetime
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

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/{types/json_variable_value.py → resources/organizations/__init__.py}

@@ -1,3 +1,3 @@
 # WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
 
-from vellum.client.types.json_variable_value import *
+from vellum.client.resources.organizations import *

vellum/resources/organizations/client.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.resources.organizations.client import *

vellum/{types/image_variable_value.py → resources/workspaces/__init__.py}

@@ -1,3 +1,3 @@
 # WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
 
-from vellum.client.types.image_variable_value import *
+from vellum.client.resources.workspaces import *

vellum/{types/array_variable_value.py → resources/workspaces/client.py}

@@ -1,3 +1,3 @@
 # WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
 
-from vellum.client.types.array_variable_value import *
+from vellum.client.resources.workspaces.client import *

vellum/types/{array_variable_value_item.py → new_member_join_behavior_enum.py}

@@ -1,3 +1,3 @@
 # WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
 
-from vellum.client.types.array_variable_value_item import *
+from vellum.client.types.new_member_join_behavior_enum import *

vellum/types/{audio_variable_value.py → organization_read.py}

@@ -1,3 +1,3 @@
 # WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
 
-from vellum.client.types.audio_variable_value import *
+from vellum.client.types.organization_read import *

vellum/types/{error_variable_value.py → workspace_read.py}

@@ -1,3 +1,3 @@
 # WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
 
-from vellum.client.types.error_variable_value import *
+from vellum.client.types.workspace_read import *

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