vellum-ai 0.12.16__tar.gz → 0.12.17__tar.gz

{vellum_ai-0.12.16 → vellum_ai-0.12.17}/PKG-INFO

@@ -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 → vellum_ai-0.12.17}/ee/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_ai-0.12.16 → vellum_ai-0.12.17}/ee/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_ai-0.12.16 → vellum_ai-0.12.17}/ee/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

{vellum_ai-0.12.16 → vellum_ai-0.12.17}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-ai"
-version = "0.12.16"
+version = "0.12.17"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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_ai-0.12.16 → vellum_ai-0.12.17}/src/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 {}