vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

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

@@ -1,12 +1,13 @@
 import io
 import os
+import sys
+import traceback
 from typing import Any, Optional, Tuple, Union
 
-from pydantic import BaseModel
-
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
-from vellum.workflows.nodes.utils import cast_to_output_type
+from vellum.workflows.nodes.utils import cast_to_output_type, wrap_inputs_for_backward_compatibility
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.core import EntityInputsInterface
 
@@ -35,62 +36,11 @@ def read_file_from_path(
         return None
 
 
-class ListWrapper(list):
-    def __getitem__(self, key):
-        item = super().__getitem__(key)
-        if not isinstance(item, DictWrapper) and not isinstance(item, ListWrapper):
-            self.__setitem__(key, _clean_for_dict_wrapper(item))
-
-        return super().__getitem__(key)
-
-
-class DictWrapper(dict):
-    """
-    This wraps a dict object to make it behave basically the same as a standard javascript object
-    and enables us to use vellum types here without a shared library since we don't actually
-    typecheck things here.
-    """
-
-    def __getitem__(self, key):
-        return self.__getattr__(key)
-
-    def __getattr__(self, attr):
-        if attr not in self:
-            if attr == "value":
-                # In order to be backwards compatible with legacy Workflows, which wrapped
-                # several values as VellumValue objects, we use the "value" key to return itself
-                return self
-
-            raise AttributeError(f"dict has no key: '{attr}'")
-
-        item = super().__getitem__(attr)
-        if not isinstance(item, DictWrapper) and not isinstance(item, ListWrapper):
-            self.__setattr__(attr, _clean_for_dict_wrapper(item))
-
-        return super().__getitem__(attr)
-
-    def __setattr__(self, name, value):
-        self[name] = value
-
-
-def _clean_for_dict_wrapper(obj):
-    if isinstance(obj, dict):
-        wrapped = DictWrapper(obj)
-        for key in wrapped:
-            wrapped[key] = _clean_for_dict_wrapper(wrapped[key])
-
-        return wrapped
-
-    elif isinstance(obj, list):
-        return ListWrapper(map(lambda item: _clean_for_dict_wrapper(item), obj))
-
-    return obj
-
-
 def run_code_inline(
     code: str,
     inputs: EntityInputsInterface,
     output_type: Any,
+    filepath: str,
 ) -> Tuple[str, Any]:
     log_buffer = io.StringIO()
 
@@ -99,38 +49,42 @@ def run_code_inline(
         print_line = f"{' '.join(str_args)}\n"
         log_buffer.write(print_line)
 
-    def wrap_value(value):
-        if isinstance(value, list):
-            return ListWrapper(
-                [
-                    # Convert VellumValue to dict with its fields
-                    (
-                        item.model_dump()
-                        if isinstance(item, BaseModel)
-                        else _clean_for_dict_wrapper(item) if isinstance(item, (dict, list)) else item
-                    )
-                    for item in value
-                ]
-            )
-        return _clean_for_dict_wrapper(value)
-
+    wrapped_inputs = wrap_inputs_for_backward_compatibility(inputs)
     exec_globals = {
-        "__arg__inputs": {name: wrap_value(value) for name, value in inputs.items()},
+        "__arg__inputs": wrapped_inputs,
         "__arg__out": None,
         "print": _inline_print,
     }
-    run_args = [f"{name}=__arg__inputs['{name}']" for name in inputs.keys()]
+    run_args = [f"{name}=__arg__inputs['{name}']" for name, value in inputs.items() if value is not undefined]
     execution_code = f"""\
 {code}
 
 __arg__out = main({", ".join(run_args)})
 """
     try:
-        exec(execution_code, exec_globals)
+        compiled_code = compile(execution_code, filepath, "exec")
+        exec(compiled_code, exec_globals)
     except Exception as e:
+        lines = code.splitlines()
+        _, _, tb = sys.exc_info()
+        tb_generator = traceback.walk_tb(tb)
+        stack = traceback.StackSummary.extract(tb_generator)
+
+        # Filter stack to only include frames from the user's code file, and omit the first one
+        filtered_stack = traceback.StackSummary.from_list([frame for frame in stack if frame.filename == filepath][1:])
+        for frame in filtered_stack:
+            if not frame.line and frame.lineno and frame.lineno <= len(lines):
+                # Mypy doesn't like us setting private attributes
+                frame._line = lines[frame.lineno - 1]  # type: ignore[attr-defined]
+
+        error_message = f"""\
+Traceback (most recent call last):
+{''.join(filtered_stack.format())}
+{e.__class__.__name__}: {e}
+"""
         raise NodeException(
             code=WorkflowErrorCode.INVALID_CODE,
-            message=str(e),
+            message=error_message,
         )
 
     logs = log_buffer.getvalue()

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

@@ -1,24 +1,28 @@
-from typing import Set
+from typing import Generic, Set
 
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
-from vellum.workflows.ports.node_ports import NodePorts
 from vellum.workflows.ports.port import Port
 from vellum.workflows.ports.utils import validate_ports
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.generics import StateType
 
 
-class ConditionalNode(BaseNode):
+class ConditionalNode(BaseNode[StateType], Generic[StateType]):
     """
     Used to conditionally determine which port to invoke next. This node exists to be backwards compatible with
     Vellum's Conditional Node, and for most cases, you should extend `BaseNode.Ports` directly.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:split"
+        color = "corn"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
-    class Ports(NodePorts):
+    class Ports(BaseNode.Ports):
         def __call__(self, outputs: BaseOutputs, state: BaseState) -> Set[Port]:
             all_ports = [port for port in self.__class__]
             enforce_single_invoked_port = validate_ports(all_ports)

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

@@ -1,9 +1,11 @@
-from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args
+from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args, get_origin
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
 from vellum.workflows.nodes.utils import cast_to_output_type
+from vellum.workflows.ports import NodePorts
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.types.utils import get_original_base
@@ -26,6 +28,7 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
             **annotations,
             "value": parent.get_output_type(),
         }
+
         return parent
 
     def get_output_type(cls) -> Type:
@@ -37,6 +40,73 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
         else:
             return all_args[1]
 
+    def __validate__(cls) -> None:
+        cls._validate_output_type_consistency(cls)
+
+    @classmethod
+    def _validate_output_type_consistency(mcs, cls: Type) -> None:
+        """
+        Validates that the declared output type of FinalOutputNode matches
+        the type of the descriptor assigned to the 'value' attribute in its Outputs class.
+
+        Raises ValueError if there's a type mismatch.
+        """
+        if not hasattr(cls, "Outputs"):
+            return
+
+        outputs_class = cls.Outputs
+        if not hasattr(outputs_class, "value"):
+            return
+
+        declared_output_type = cls.get_output_type()
+
+        if declared_output_type is Any:
+            return
+
+        value_descriptor = None
+
+        if "value" in outputs_class.__dict__:
+            value_descriptor = outputs_class.__dict__["value"]
+        else:
+            value_descriptor = getattr(outputs_class, "value")
+
+        if isinstance(value_descriptor, OutputReference):
+            descriptor_types = value_descriptor.types
+
+            type_mismatch = True
+            for descriptor_type in descriptor_types:
+                if descriptor_type == declared_output_type:
+                    type_mismatch = False
+                    break
+                if (
+                    get_origin(descriptor_type) == declared_output_type
+                    or get_origin(declared_output_type) == descriptor_type
+                ):
+                    type_mismatch = False
+                    break
+                try:
+                    if issubclass(descriptor_type, declared_output_type) or issubclass(
+                        declared_output_type, descriptor_type
+                    ):
+                        type_mismatch = False
+                        break
+                except TypeError:
+                    # Handle cases where types aren't classes (e.g., Union)
+                    if str(descriptor_type) == str(declared_output_type):
+                        type_mismatch = False
+                        break
+
+            if type_mismatch:
+                declared_type_name = getattr(declared_output_type, "__name__", str(declared_output_type))
+                descriptor_type_names = [getattr(t, "__name__", str(t)) for t in descriptor_types]
+
+                raise ValueError(
+                    f"Output type mismatch in {cls.__name__}: "
+                    f"FinalOutputNode is declared with output type '{declared_type_name}' "
+                    f"but the 'value' descriptor has type(s) {descriptor_type_names}. "
+                    f"The output descriptor type must match the declared FinalOutputNode output type."
+                )
+
 
 class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], metaclass=_FinalOutputNodeMeta):
     """
@@ -44,9 +114,16 @@ class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], meta
     This provides backward compatibility with Vellum's Final Output Node.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:circle-stop"
+        color = "teal"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
+    class Ports(NodePorts):
+        pass
+
     class Outputs(BaseNode.Outputs):
         # We use our mypy plugin to override the _OutputType with the actual output type
         # for downstream references to this output.
@@ -61,3 +138,5 @@ class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], meta
                 self.__class__.get_output_type(),
             )
         )
+
+    __simulates_workflow_output__ = True

vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py

@@ -1,11 +1,16 @@
 import pytest
+from typing import Any, Dict
 
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.nodes.displayable.inline_prompt_node import InlinePromptNode
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import Json
 
 
-def test_final_output_node__mismatched_output_type():
+def test_final_output_node__mismatched_output_type_should_raise_exception_when_ran():
     # GIVEN a FinalOutputNode with a mismatched output type
     class StringOutputNode(FinalOutputNode[BaseState, str]):
         class Outputs(FinalOutputNode.Outputs):
@@ -18,3 +23,93 @@ def test_final_output_node__mismatched_output_type():
 
     # THEN an error is raised
     assert str(exc_info.value) == "Expected an output of type 'str', but received 'dict'"
+
+
+def test_final_output_node__mismatched_output_type_should_raise_exception():
+    # GIVEN a FinalOutputNode declared with list output type but has a string value type
+    class Output(FinalOutputNode[BaseState, list]):
+        """Output the extracted invoice line items as an array of objects."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = InlinePromptNode.Outputs.text
+
+    # WHEN attempting to validate the node class
+    # THEN a ValueError should be raised during validation
+    with pytest.raises(ValueError) as exc_info:
+        Output.__validate__()
+
+    # AND the error message should indicate the type mismatch
+    assert (
+        str(exc_info.value)
+        == "Output type mismatch in Output: FinalOutputNode is declared with output type 'list' but "
+        "the 'value' descriptor has type(s) ['str']. The output descriptor type must match the "
+        "declared FinalOutputNode output type."
+    )
+
+
+def test_final_output_node__matching_output_type_should_pass_validation():
+    # GIVEN a FinalOutputNode declared with correct matching types
+    class CorrectOutput(FinalOutputNode[BaseState, str]):
+        """Output with correct type matching."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = InlinePromptNode.Outputs.text
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        CorrectOutput.__validate__()
+    except ValueError:
+        pytest.fail("Validation should not raise an exception for correct type matching")
+
+
+def test_final_output_node__dict_and_Dict_should_be_compatible():
+    """
+    Tests that FinalOutputNode validation recognizes dict and Dict[str, Any] as compatible types.
+    """
+
+    # GIVEN a FinalOutputNode declared with dict output type
+    # AND the value descriptor has Dict[str, Any] type
+    class DictOutputNode(FinalOutputNode[BaseState, dict]):
+        """Output with dict type."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = OutputReference(
+                name="value",
+                types=(Dict[str, Any],),
+                instance=None,
+                outputs_class=FinalOutputNode.Outputs,
+            )
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        DictOutputNode.__validate__()
+    except ValueError as e:
+        pytest.fail(f"Validation should not raise an exception for dict/Dict compatibility: {e}")
+
+
+def test_final_output_node__any_output_type_should_accept_json():
+    """
+    Tests that FinalOutputNode with Any output type accepts a TemplatingNode with Json output type.
+    """
+
+    # GIVEN a TemplatingNode with Json output type
+    class JsonTemplatingNode(TemplatingNode[BaseState, Json]):
+        """Templating node that outputs Json."""
+
+        template = '{"key": "value"}'
+
+    # AND a FinalOutputNode with Any output type referencing the TemplatingNode
+    class AnyOutputNode(FinalOutputNode[BaseState, Any]):
+        """Output with Any type."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = JsonTemplatingNode.Outputs.result
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        AnyOutputNode.__validate__()
+    except ValueError as e:
+        pytest.fail(f"Validation should not raise an exception when Any accepts Json: {e}")

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

@@ -3,12 +3,11 @@ from typing import Any, ClassVar, Dict, Generic, List, Optional, Union, cast
 
 from vellum import ChatHistoryInput, ChatMessage, JsonInput, MetricDefinitionInput, NumberInput, StringInput
 from vellum.client import ApiError
-from vellum.core import RequestOptions
+from vellum.client.core import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.generics import StateType
 
@@ -23,6 +22,10 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
     request_options: Optional[RequestOptions] - The request options to use for the Metric
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:shield-check"
+        color = "rose"
+
     metric_definition: ClassVar[Union[UUID, str]]
 
     metric_inputs: ClassVar[EntityInputsInterface] = {}
@@ -33,10 +36,11 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
-    class Outputs(BaseOutputs):
+    class Outputs(BaseNode.Outputs):
         score: float
         normalized_score: Optional[float]
         log: Optional[str]
+        reason: Optional[str]
 
     def run(self) -> Outputs:
         try:
@@ -54,10 +58,10 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
             )
 
         metric_outputs = {output.name: output.value for output in metric_execution.outputs}
-
         SCORE_KEY = "score"
         NORMALIZED_SCORE_KEY = "normalized_score"
         LOG_KEY = "log"
+        REASON_KEY = "reason"
 
         score = metric_outputs.get(SCORE_KEY)
         if not isinstance(score, float):
@@ -87,7 +91,23 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
         else:
             log = None
 
-        return self.Outputs(score=score, normalized_score=normalized_score, log=log, **metric_outputs)
+        if REASON_KEY in metric_outputs:
+            reason = metric_outputs.pop(REASON_KEY) or ""
+            if not isinstance(reason, str):
+                raise NodeException(
+                    message="Metric execution reason output must be of type 'str'",
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
+                )
+        else:
+            reason = None
+
+        return self.Outputs(
+            score=score,
+            normalized_score=normalized_score,
+            log=log,
+            reason=reason,
+            **metric_outputs,  # type: ignore [arg-type]
+        )
 
     def _compile_metric_inputs(self) -> List[MetricDefinitionInput]:
         # TODO: We may want to consolidate with prompt deployment input compilation

vellum/workflows/nodes/displayable/guardrail_node/test_node.py

@@ -102,6 +102,35 @@ def test_run_guardrail_node__normalized_score_null(vellum_client):
     assert exc_info.value.code == WorkflowErrorCode.INVALID_OUTPUTS
 
 
+def test_run_guardrail_node__reason(vellum_client):
+    # GIVEN a Guardrail Node
+    class MyGuard(GuardrailNode):
+        metric_definition = "example_metric_definition"
+        metric_inputs = {}
+
+    # AND we know that the guardrail node will return a reason
+    mock_metric_execution = MetricDefinitionExecution(
+        outputs=[
+            TestSuiteRunMetricNumberOutput(
+                name="score",
+                value=0.6,
+            ),
+            TestSuiteRunMetricStringOutput(
+                name="reason",
+                value="foo",
+            ),
+        ],
+    )
+    vellum_client.metric_definitions.execute_metric_definition.return_value = mock_metric_execution
+
+    # WHEN we run the Guardrail Node
+    outputs = MyGuard().run()
+
+    # THEN the workflow should have completed successfully
+    assert outputs.score == 0.6
+    assert outputs.reason == "foo"
+
+
 def test_run_guardrail_node__api_error(vellum_client):
     # GIVEN a Guardrail Node
     class MyGuard(GuardrailNode):

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

@@ -1,16 +1,16 @@
-import json
-from typing import Any, Dict, Iterator, Type, Union
+from typing import Any, Dict, Generic, Iterator, Type, Union
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BaseInlinePromptNode as BaseInlinePromptNode
+from vellum.workflows.nodes.displayable.bases.utils import process_additional_prompt_outputs
 from vellum.workflows.outputs import BaseOutput
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 
 
-class InlinePromptNode(BaseInlinePromptNode[StateType]):
+class InlinePromptNode(BaseInlinePromptNode[StateType], Generic[StateType]):
     """
     Used to execute a Prompt defined inline.
 
@@ -23,6 +23,10 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
     request_options: Optional[RequestOptions] - The request options to use for the Prompt Execution
     """
 
+    class Display(BaseInlinePromptNode.Display):
+        icon = "vellum:icon:text-size"
+        color = "navy"
+
     class Trigger(BaseInlinePromptNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
@@ -45,39 +49,8 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
                 code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
-        string_outputs = []
-        json_output = None
-
-        should_parse_json = False
-        if hasattr(self, "parameters"):
-            custom_params = self.parameters.custom_parameters
-            if custom_params and isinstance(custom_params, dict):
-                json_schema = custom_params.get("json_schema", {})
-                if (isinstance(json_schema, dict) and "schema" in json_schema) or custom_params.get("json_mode", {}):
-                    should_parse_json = True
-
-        for output in outputs:
-            if output.value is None:
-                continue
-
-            if output.type == "STRING":
-                string_outputs.append(output.value)
-                if should_parse_json:
-                    try:
-                        parsed_json = json.loads(output.value)
-                        json_output = parsed_json
-                    except (json.JSONDecodeError, TypeError):
-                        pass
-            elif output.type == "JSON":
-                string_outputs.append(json.dumps(output.value, indent=4))
-                json_output = output.value
-            elif output.type == "FUNCTION_CALL":
-                string_outputs.append(output.value.model_dump_json(indent=4))
-            else:
-                string_outputs.append(output.value.message)
-
-        value = "\n".join(string_outputs)
-        yield BaseOutput(name="text", value=value)
+        text_output, json_output = process_additional_prompt_outputs(outputs)
+        yield BaseOutput(name="text", value=text_output)
 
         if json_output:
             yield BaseOutput(name="json", value=json_output)

vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py

@@ -197,6 +197,9 @@ def test_inline_prompt_node__api_error__invalid_inputs_node_exception(
     assert e.value.code == expected_code
     assert e.value.message == expected_message
 
+    # AND the node exception includes the response body in raw_data
+    assert e.value.raw_data == exception.body
+
 
 def test_inline_prompt_node__chat_history_inputs(vellum_adhoc_prompt_client):
     # GIVEN a prompt node with a chat history input

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

@@ -1,12 +1,19 @@
+from typing import Generic
+
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.generics import StateType
 
 
-class MergeNode(BaseNode):
+class MergeNode(BaseNode[StateType], Generic[StateType]):
     """
     Used to merge the control flow of multiple nodes into a single node. This node exists to be backwards compatible
     with Vellum's Merge Node, and for most cases, you should extend from `BaseNode.Trigger` directly.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:merge"
+        color = "tomato"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY

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

@@ -1,12 +1,19 @@
+from typing import Generic
+
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.generics import StateType
 
 
-class NoteNode(BaseNode):
+class NoteNode(BaseNode[StateType], Generic[StateType]):
     """
     A no-op Node purely used to display a note in the Vellum UI.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:note"
+        color = "cyan"
+
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY