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

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

@@ -1,7 +1,9 @@
 import enum
 import json
-from typing import Any, List, Union, cast
+import re
+from typing import Any, List, Optional, Tuple, Union, cast
 
+from vellum import PromptOutput
 from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.client.types.array_vellum_value_request import ArrayVellumValueRequest
 from vellum.client.types.audio_vellum_value import AudioVellumValue
@@ -28,15 +30,18 @@ from vellum.client.types.string_vellum_value_request import StringVellumValueReq
 from vellum.client.types.vellum_error import VellumError
 from vellum.client.types.vellum_value import VellumValue
 from vellum.client.types.vellum_value_request import VellumValueRequest
+from vellum.client.types.video_vellum_value import VideoVellumValue
+from vellum.client.types.video_vellum_value_request import VideoVellumValueRequest
+from vellum.utils.json_encoder import VellumJsonEncoder
 from vellum.workflows.errors.types import WorkflowError, workflow_error_to_vellum_error
-from vellum.workflows.state.encoder import DefaultStateEncoder
 
 VELLUM_VALUE_REQUEST_TUPLE = (
     StringVellumValueRequest,
     NumberVellumValueRequest,
     JsonVellumValueRequest,
-    ImageVellumValueRequest,
     AudioVellumValueRequest,
+    VideoVellumValueRequest,
+    ImageVellumValueRequest,
     FunctionCallVellumValueRequest,
     ErrorVellumValueRequest,
     ArrayVellumValueRequest,
@@ -52,6 +57,8 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
         return StringVellumValue(value=value)
     elif isinstance(value, enum.Enum):
         return StringVellumValue(value=value.value)
+    elif isinstance(value, bool):
+        return JsonVellumValue(value=value)
     elif isinstance(value, (int, float)):
         return NumberVellumValue(value=value)
     elif isinstance(value, list) and (
@@ -76,8 +83,9 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
             StringVellumValue,
             NumberVellumValue,
             JsonVellumValue,
-            ImageVellumValue,
             AudioVellumValue,
+            VideoVellumValue,
+            ImageVellumValue,
             FunctionCallVellumValue,
             ErrorVellumValue,
             ArrayVellumValue,
@@ -95,7 +103,7 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
         return value  # type: ignore
 
     try:
-        json_value = json.dumps(value, cls=DefaultStateEncoder)
+        json_value = json.dumps(value, cls=VellumJsonEncoder)
     except json.JSONDecodeError:
         raise ValueError(f"Unsupported variable type: {value.__class__.__name__}")
 
@@ -117,3 +125,72 @@ def primitive_to_vellum_value_request(value: Any) -> VellumValueRequest:
         raise ValueError(f"Unsupported variable type: {vellum_value.__class__.__name__}")
 
     return vellum_value_request_class.model_validate(vellum_value.model_dump())
+
+
+def _strip_markdown_json(text: str) -> str:
+    """
+    Strip markdown code block formatting from JSON text.
+
+    Handles cases like:
+    ```json
+    {"key": "value"}
+    ```
+
+    Or with surrounding text:
+    Here is the json you asked for!
+
+    ```json
+    {"key": "value"}
+    ```
+
+    Args:
+        text: The text potentially containing markdown-formatted JSON
+
+    Returns:
+        The stripped JSON text, or the original text if no markdown formatting found
+    """
+    pattern = r"```(?:json)?\s*\n(.*?)\n```"
+
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+
+    return text
+
+
+def process_additional_prompt_outputs(outputs: List[PromptOutput]) -> Tuple[str, Optional[Any]]:
+    """
+    Process prompt outputs using the same logic as prompt nodes to determine text output.
+
+    Args:
+        outputs: List of prompt outputs to process
+
+    Returns:
+        The text representation of the outputs joined with newlines
+        The JSON representation of the outputs
+    """
+    string_outputs = []
+    json_output = None
+    for output in outputs:
+        if output.value is None:
+            continue
+
+        if output.type == "STRING":
+            string_outputs.append(output.value)
+            for value_to_parse in [output.value, _strip_markdown_json(output.value)]:
+                try:
+                    json_output = json.loads(value_to_parse)
+                    break
+                except (json.JSONDecodeError, TypeError):
+                    continue
+        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))
+        elif output.type == "THINKING":
+            continue
+        else:
+            string_outputs.append(output.value.message)
+
+    return "\n".join(string_outputs), json_output

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

@@ -20,9 +20,10 @@ from vellum import (
     VellumError,
     VellumValue,
 )
+from vellum.client.core import RequestOptions
 from vellum.client.core.api_error import ApiError
 from vellum.client.types.code_executor_secret_input import CodeExecutorSecretInput
-from vellum.core import RequestOptions
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
@@ -77,6 +78,10 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
     request_options: Optional[RequestOptions] = None - The request options to use for the custom script.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:rectangle-code"
+        color = "lime"
+
     filepath: ClassVar[Optional[str]] = None
     code: ClassVar[Optional[str]] = None
 
@@ -98,7 +103,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
     def run(self) -> Outputs:
         output_type = self.__class__.get_output_type()
         code, filepath = self._resolve_code()
-        if not self.packages and self.runtime == "PYTHON_3_11_6":
+        if not self.packages and self.runtime == "PYTHON_3_11_6" and not self._has_secrets_in_code_inputs():
             logs, result = run_code_inline(code, self.code_inputs, output_type, filepath)
             return self.Outputs(result=result, log=logs)
 
@@ -116,13 +121,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                     request_options=self.request_options,
                 )
             except ApiError as e:
-                if e.status_code == 400 and isinstance(e.body, dict) and "message" in e.body:
-                    raise NodeException(
-                        message=e.body["message"],
-                        code=WorkflowErrorCode.INVALID_INPUTS,
-                    )
-
-                raise
+                self._handle_api_error(e)
 
             if code_execution_result.output.type != expected_output_type:
                 actual_type = code_execution_result.output.type
@@ -133,6 +132,31 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
             return self.Outputs(result=code_execution_result.output.value, log=code_execution_result.log)
 
+    def _handle_api_error(self, e: ApiError) -> None:
+        if e.status_code and e.status_code == 403 and isinstance(e.body, dict):
+            raise NodeException(
+                message=e.body.get("detail", "Provider credentials is missing or unavailable"),
+                code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+            ) from e
+
+        if e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
+            raise NodeException(
+                message=e.body.get("message", e.body.get("detail", "Failed to execute code")),
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            ) from e
+
+        raise NodeException(
+            message="Failed to execute code",
+            code=WorkflowErrorCode.INTERNAL_ERROR,
+        ) from e
+
+    def _has_secrets_in_code_inputs(self) -> bool:
+        """Check if any code_inputs contain VellumSecret instances that require API execution."""
+        for input_value in self.code_inputs.values():
+            if isinstance(input_value, VellumSecret):
+                return True
+        return False
+
     def _compile_code_inputs(self) -> List[CodeExecutorInput]:
         # TODO: We may want to consolidate with prompt deployment input compilation
         # https://app.shortcut.com/vellum/story/4117
@@ -140,6 +164,8 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
         compiled_inputs: List[CodeExecutorInput] = []
 
         for input_name, input_value in self.code_inputs.items():
+            if input_value is undefined:
+                continue
             if isinstance(input_value, str):
                 compiled_inputs.append(
                     StringInput(

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

@@ -6,13 +6,18 @@ from typing import Any, List, Union
 from pydantic import BaseModel
 
 from vellum import ArrayInput, CodeExecutorResponse, NumberVellumValue, StringInput, StringVellumValue
+from vellum.client.core.api_error import ApiError
 from vellum.client.errors.bad_request_error import BadRequestError
+from vellum.client.errors.forbidden_error import ForbiddenError
+from vellum.client.errors.internal_server_error import InternalServerError
+from vellum.client.errors.not_found_error import NotFoundError
 from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.code_execution_package import CodeExecutionPackage
 from vellum.client.types.code_executor_secret_input import CodeExecutorSecretInput
 from vellum.client.types.function_call import FunctionCall
 from vellum.client.types.number_input import NumberInput
 from vellum.client.types.string_chat_message_content import StringChatMessageContent
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
@@ -109,6 +114,7 @@ def main(word: str) -> int:
             CodeExecutionPackage(
                 name="openai",
                 version="1.0.0",
+                repository="test-repo",
             )
         ]
 
@@ -150,6 +156,7 @@ def main(word: str) -> int:
             CodeExecutionPackage(
                 name="openai",
                 version="1.0.0",
+                repository="test-repo",
             )
         ],
         request_options=None,
@@ -789,6 +796,103 @@ Node.js v21.7.3
     assert exc_info.value.message == message
 
 
+def test_run_node__execute_code_api_fails_403__provider_credentials_unavailable(vellum_client):
+    """
+    Tests that a 403 error from the API is handled with PROVIDER_CREDENTIALS_UNAVAILABLE error code.
+    """
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = "def main(): return 'test'"
+        runtime = "PYTHON_3_11_6"
+        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]
+
+    vellum_client.execute_code.side_effect = ForbiddenError(
+        body={
+            "detail": "Provider credentials is missing or unavailable",
+        }
+    )
+
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE
+    assert "Provider credentials is missing or unavailable" in exc_info.value.message
+
+
+def test_run_node__execute_code_api_fails_404__invalid_inputs(vellum_client):
+    """
+    Tests that a 404 error from the API is handled with INVALID_INPUTS error code.
+    """
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = "def main(): return 'test'"
+        runtime = "PYTHON_3_11_6"
+        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]
+
+    vellum_client.execute_code.side_effect = NotFoundError(
+        body={
+            "detail": "Resource not found",
+        }
+    )
+
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "Resource not found" in exc_info.value.message
+
+
+def test_run_node__execute_code_api_fails_500__internal_error(vellum_client):
+    """
+    Tests that a 500 error from the API is handled with INTERNAL_ERROR error code.
+    """
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = "def main(): return 'test'"
+        runtime = "PYTHON_3_11_6"
+        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]
+
+    vellum_client.execute_code.side_effect = InternalServerError(
+        body={
+            "detail": "Internal server error occurred",
+        }
+    )
+
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INTERNAL_ERROR
+    assert exc_info.value.message == "Failed to execute code"
+
+
+def test_run_node__execute_code_api_fails_4xx_no_message__uses_detail(vellum_client):
+    """
+    Tests that a 4xx error without a 'message' field falls back to 'detail' field.
+    """
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = "def main(): return 'test'"
+        runtime = "PYTHON_3_11_6"
+        packages = [CodeExecutionPackage(name="requests", version="2.28.0")]
+
+    vellum_client.execute_code.side_effect = ApiError(
+        status_code=422,
+        body={
+            "detail": "Invalid request parameters",
+        },
+    )
+
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "Invalid request parameters" in exc_info.value.message
+
+
 def test_run_node__execute_code__list_extends():
     # GIVEN a node that will return a list with output type Json
     class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Json]):
@@ -845,21 +949,11 @@ def main(arg1: list) -> str:
         runtime = "PYTHON_3_11_6"
 
     # WHEN we run the node
-    with pytest.raises(NodeException) as exc_info:
-        node = ExampleCodeExecutionNode()
-        node.run()
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
 
     # AND the result should be the correct output
-    assert (
-        exc_info.value.message
-        == """\
-Traceback (most recent call last):
-  File "ExampleCodeExecutionNode.code.py", line 2, in main
-    return arg1["invalid"]
-
-AttributeError: dict has no key: 'invalid'
-"""
-    )
+    assert outputs == {"result": "", "log": ""}
 
 
 def test_run_node__execute_code__value_key_access():
@@ -1233,3 +1327,138 @@ def main(numbers: List[str]) -> str:
         ],
         request_options=None,
     )
+
+
+def test_run_node__vellum_secret_forces_api_execution(vellum_client):
+    """Test that CodeExecutionNode with VellumSecretReference forces API execution instead of inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode with VellumSecretReference but no packages
+    # This should normally run inline, but the presence of secrets should force API execution
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, str]):
+        code = """
+def main(secret: str) -> str:
+    return f"Secret value: {secret}"
+"""
+        runtime = "PYTHON_3_11_6"
+        # Note: No packages specified, which would normally trigger inline execution
+
+        code_inputs = {
+            "secret": VellumSecretReference("MY_SECRET"),
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="",
+        output=StringVellumValue(value="Secret value: my_secret_value"),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": "Secret value: my_secret_value", "log": ""}
+
+    # AND we should have invoked the Code via API (not inline) due to the secret
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            CodeExecutorSecretInput(
+                name="secret",
+                value="MY_SECRET",
+            )
+        ],
+        code="""
+def main(secret: str) -> str:
+    return f"Secret value: {secret}"
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="STRING",
+        packages=[],
+        request_options=None,
+    )
+
+
+def test_run_node__undefined_input_skipped():
+    """
+    Confirm that when an undefined value is passed as an input, it is skipped rather than raising an error.
+    The function should use the default parameter value when undefined input is skipped.
+    """
+
+    # GIVEN a node with both a valid input and an undefined input that runs inline
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        code = """\
+def main(word: str, undefined_input: str = "default") -> int:
+    return len(word) + len(undefined_input)
+"""
+        runtime = "PYTHON_3_11_6"
+        packages = []
+
+        code_inputs = {
+            "word": "hello",
+            "undefined_input": undefined,
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+    outputs = node.run()
+
+    # THEN the node should run successfully without raising an error
+    assert outputs == {"result": 12, "log": ""}  # len("hello") + len("default") = 5 + 7 = 12
+
+
+def test_run_node__undefined_input_skipped_api_execution(vellum_client):
+    """
+    Confirm that when an undefined value is passed as an input in API execution mode,
+    it is skipped rather than raising an error.
+    """
+
+    # GIVEN a node with both a valid input and an undefined input that forces API execution
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        code = """\
+def main(word: str) -> int:
+    return len(word)
+"""
+        runtime = "PYTHON_3_11_6"
+        packages = [CodeExecutionPackage(name="requests", version="2.0.0")]
+
+        code_inputs = {
+            "word": "hello",
+            "undefined_input": undefined,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+    outputs = node.run()
+
+    # THEN the node should run successfully without raising an error
+    assert outputs == {"result": 5, "log": ""}
+
+    # AND the API should be called with only the non-undefined input
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[StringInput(name="word", value="hello")],
+        code="""\
+def main(word: str) -> int:
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[CodeExecutionPackage(name="requests", version="2.0.0")],
+        request_options=None,
+    )

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

@@ -4,13 +4,11 @@ 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.code_execution_node_wrappers import ListWrapper, clean_for_dict_wrapper
 from vellum.workflows.types.core import EntityInputsInterface
 
 
@@ -51,27 +49,13 @@ 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, str)) 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}
 

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)