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

vellum/workflows/nodes/displayable/code_execution_node/tests/{test_code_execution_node.py → test_node.py}

@@ -5,14 +5,19 @@ from typing import Any, List, Union
 
 from pydantic import BaseModel
 
-from vellum import CodeExecutorResponse, NumberVellumValue, StringInput, StringVellumValue
+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,
@@ -409,6 +416,56 @@ def main(word: str) -> int:
     assert exc_info.value.message == "Expected an output of type 'int', but received 'str'"
 
 
+def test_run_node__run_inline__incorrect_output_type_list():
+    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode that returns a list but is defined to return an int
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
+        code = """\
+def main(output: list[str]) -> int:
+    return output.value
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "output": ["hello", "world"],
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Expected an output of type 'int', but received 'list'"
+
+
+def test_run_node__run_inline__incorrect_output_type_dict():
+    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode that returns a dict but is defined to return an int
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
+        code = """\
+def main(output: dict[str, str]) -> int:
+    return output.value
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "output": {
+                "foo": "bar",
+            }
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Expected an output of type 'int', but received 'dict'"
+
+
 def test_run_node__run_inline__valid_dict_to_pydantic():
     """Confirm that CodeExecutionNodes can convert a dict to a Pydantic model during inline execution."""
 
@@ -739,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]):
@@ -795,12 +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 == "dict has no key: 'invalid'"
+    assert outputs == {"result": "", "log": ""}
 
 
 def test_run_node__execute_code__value_key_access():
@@ -821,3 +974,491 @@ def main(arg1: list) -> str:
 
     # AND the result should be the correct output
     assert outputs == {"result": "bar", "log": ""}
+
+
+@pytest.mark.parametrize(
+    "code_snippet",
+    [
+        """
+def main(text: str) -> str:
+    return text.value
+""",
+        """
+def main(text: str) -> str:
+    return text["value"]
+""",
+    ],
+)
+def test_run_node__string_value_wrapper_value(code_snippet):
+    """Test string value wrapper value access using different patterns"""
+
+    # GIVEN a node that accesses the 'value' property of a string input
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = code_snippet
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "text": "hello",
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value
+    assert outputs == {"result": "hello", "log": ""}
+
+
+@pytest.mark.parametrize(
+    "code_snippet",
+    [
+        """
+def main(text: str) -> str:
+    return text.type
+""",
+        """
+def main(text: str) -> str:
+    return text["type"]
+""",
+    ],
+)
+def test_run_node__string_value_wrapper_type(code_snippet):
+    """Test string value wrapper type access using different patterns"""
+
+    # GIVEN a node that will return the string type
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = code_snippet
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "text": "hello",
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully return the string type
+    assert outputs == {"result": "STRING", "log": ""}
+
+
+def test_run_node__string_value_wrapper__list_of_dicts():
+    # GIVEN a node that accesses the 'value' property of a string input
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
+        code = """\
+def main(output: list[str]) -> list[str]:
+    results = []
+    for item in output:
+        results.append(item['value'])
+
+    return results
+"""
+        code_inputs = {"output": ['{"foo": "bar"}', '{"foo2": "bar2"}']}
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value
+    assert outputs == {"result": ['{"foo": "bar"}', '{"foo2": "bar2"}'], "log": ""}
+
+
+def test_run_node__string_key_access_still_works():
+    # GIVEN a node that accesses the '0' index of a string input
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
+        code = """\
+def main(input: str) -> str:
+    return input[0]
+"""
+        code_inputs = {"input": "hello"}
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value
+    assert outputs == {"result": "h", "log": ""}
+
+
+@pytest.mark.parametrize(
+    "code_snippet",
+    [
+        """
+from vellum.client.types.function_call import FunctionCall
+def main(input: FunctionCall) -> FunctionCall:
+    return input.value
+""",
+        """
+from vellum.client.types.function_call import FunctionCall
+def main(input: FunctionCall) -> FunctionCall:
+    return input["value"]
+""",
+    ],
+)
+def test_run_node__function_call_wrapper_value(code_snippet):
+    """Test function call wrapper value access using different patterns"""
+
+    # GIVEN a node that accesses the function call value
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, FunctionCall]):
+        code = code_snippet
+        code_inputs = {
+            "input": FunctionCall(
+                name="test-name",
+                arguments={
+                    "test-key": "test-value",
+                },
+            )
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully return the function call value
+    assert isinstance(outputs.result, FunctionCall)
+    assert outputs.result.name == "test-name"
+    assert outputs.result.arguments == {"test-key": "test-value"}
+    assert outputs.result.id is None
+    assert outputs.log == ""
+
+
+@pytest.mark.parametrize(
+    "code_snippet",
+    [
+        """
+from vellum.client.types.function_call import FunctionCall
+def main(input: FunctionCall) -> str:
+    return input.type
+""",
+        """
+from vellum.client.types.function_call import FunctionCall
+def main(input: FunctionCall) -> str:
+    return input["type"]
+""",
+    ],
+)
+def test_run_node__function_call_wrapper_type(code_snippet):
+    """Test function call wrapper type access using different patterns"""
+
+    # GIVEN a node that accesses the function call type
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = code_snippet
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "input": FunctionCall(
+                name="test-name",
+                arguments={
+                    "test-key": "test-value",
+                },
+            )
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully return the function call type
+    assert outputs == {"result": "FUNCTION_CALL", "log": ""}
+
+
+def test_run_node__iter_list():
+    # GIVEN a node that will return the first string in a list
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = """\
+def main(
+    input_list: list
+) -> str:
+    return next((
+        o.value for o in input_list if o.type == "STRING"
+    ), None)
+"""
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "input_list": [
+                StringVellumValue(value="foo"),
+                NumberVellumValue(value=1),
+            ]
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value
+    assert outputs == {"result": "foo", "log": ""}
+
+
+def test_run_node__iter_dict():
+    # GIVEN a node that will return the first string in a list
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, list[str]]):
+        code = """\
+def main(input_dict: dict) -> list[str]:
+    return [item.value for item in input_dict.values()]
+"""
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {
+            "input_dict": {
+                "foo": "bar",
+                "baz": "qux",
+            }
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value
+    assert outputs == {"result": ["bar", "qux"], "log": ""}
+
+
+def test_run_node__show_clearer_runtime_error_message():
+    # GIVEN a node that will return the first string in a list
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, list[str]]):
+        code = """\
+def main(items: list[str]) -> list[str]:
+    return first(items)
+
+# Helper function to prove out stack traces
+def first(items: list[str]) -> str:
+    return items[0]
+"""
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {"items": []}
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should successfully access the string value
+    assert (
+        exc_info.value.message
+        == """\
+Traceback (most recent call last):
+  File "ExampleCodeExecutionNode.code.py", line 2, in main
+    return first(items)
+  File "ExampleCodeExecutionNode.code.py", line 6, in first
+    return items[0]
+
+IndexError: list index out of range
+"""
+    )
+
+
+def test_run_node__default_function_call_type():
+    # GIVEN a node that will return a FunctionCall
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, FunctionCall]):
+        code = """\
+def main(input: str) -> str:
+    return None
+"""
+        runtime = "PYTHON_3_11_6"
+        code_inputs = {"input": "foo"}
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should return default function call
+    assert outputs == {"result": FunctionCall(name="", arguments={}), "log": ""}
+
+
+def test_run_node__array_input_with_primitive_values(vellum_client):
+    # GIVEN a node that subclasses CodeExecutionNode that processes an array of primitive values
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, str]):
+        code = """\
+from typing import List
+def main(numbers: List[str]) -> str:
+    return " ".join(numbers)
+"""
+        runtime = "PYTHON_3_11_6"
+        packages = [
+            CodeExecutionPackage(
+                name="openai",
+                version="1.0.0",
+            )
+        ]
+
+        code_inputs = {
+            "numbers": ["6", "4", "2"],
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="",
+        output=StringVellumValue(value="6 4 2"),
+    )
+    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": "6 4 2", "log": ""}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            ArrayInput(
+                name="numbers",
+                value=[
+                    StringVellumValue(value="6"),
+                    StringVellumValue(value="4"),
+                    StringVellumValue(value="2"),
+                ],
+            )
+        ],
+        code="""\
+from typing import List
+def main(numbers: List[str]) -> str:
+    return " ".join(numbers)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="STRING",
+        packages=[
+            CodeExecutionPackage(
+                name="openai",
+                version="1.0.0",
+            )
+        ],
+        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,
+    )