vellum-ai 0.14.48__tar.gz → 0.14.49__tar.gz

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.48
+Version: 0.14.49
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -79,8 +79,8 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
 
         packages = raise_if_descriptor(node.packages)
 
-        _, output_display = display_context.global_node_output_displays[node.Outputs.result]
-        _, log_output_display = display_context.global_node_output_displays[node.Outputs.log]
+        output_display = self.output_display[node.Outputs.result]
+        log_output_display = self.output_display[node.Outputs.log]
 
         output_type = primitive_type_to_vellum_variable_type(node.get_output_type())
 

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py

@@ -2,7 +2,9 @@ import pytest
 from uuid import UUID
 from typing import Type
 
+from vellum.client.core.api_error import ApiError
 from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
+from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.vellum.code_execution_node import BaseCodeExecutionNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -110,3 +112,44 @@ def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
             },
         },
     ]
+
+
+def test_serialize_node__with_unresolved_secret_references(vellum_client):
+    # GIVEN a node has access to a secret reference
+    class MyNode(CodeExecutionNode):
+        code_inputs = {"api_key": VellumSecretReference("MY_API_KEY")}
+
+    # AND the secret is not found
+    vellum_client.workspace_secrets.retrieve.side_effect = ApiError(
+        status_code=404,
+        body="Secret not found",
+    )
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    data: dict = workflow_display.serialize()
+
+    # THEN the condition should be serialized correctly
+    node = next(node for node in data["workflow_raw_data"]["nodes"] if node["type"] == "CODE_EXECUTION")
+    assert node["inputs"][0]["value"] == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "type": "WORKSPACE_SECRET",
+                "data": {
+                    "type": "STRING",
+                    "workspace_secret_id": None,
+                },
+            }
+        ],
+    }
+
+    # AND we should have a warning of the invalid reference
+    # TODO: Come up with a proposal for how nodes should propagate warnings
+    # warnings = list(workflow_display.errors)
+    # assert len(warnings) == 1
+    # assert "Failed to resolve secret reference 'MY_API_KEY'" in str(warnings[0])

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -1,4 +1,3 @@
-import pytest
 from uuid import uuid4
 from typing import List
 
@@ -257,11 +256,58 @@ def test_serialize_node__workflow_input_as_nested_chat_history():
 
     # WHEN the workflow is serialized
     workflow_display = get_workflow_display(workflow_class=Workflow)
-    with pytest.raises(Exception) as exc_info:
-        workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the attribute reference
+    generic_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(GenericNode.__id__)
+    )
 
-    # THEN we should raise a user facing error
-    assert str(exc_info.value) == "Failed to serialize attribute 'attr': Nested references are not supported."
+    assert not DeepDiff(
+        {
+            "id": "11be9d37-0069-4695-a317-14a3b6519d4e",
+            "label": "test_serialize_node__workflow_input_as_nested_chat_history.<locals>.GenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "GenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "d4548468-85a4-449e-92e0-e2d8b8fd097c", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "c4a9a57d-1380-4689-8500-e8a0b6477291", "name": "default", "type": "DEFAULT"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "e878bbc9-1231-461e-9e9d-947604da116e",
+                    "name": "attr",
+                    "value": {
+                        "type": "DICTIONARY_REFERENCE",
+                        "entries": [
+                            {
+                                "key": "hello",
+                                "value": {
+                                    "type": "WORKFLOW_INPUT",
+                                    "input_variable_id": "f727c3f9-f27f-4ac9-abd7-12bf612a094e",
+                                },
+                            }
+                        ],
+                    },
+                }
+            ],
+            "outputs": [],
+        },
+        generic_node,
+        ignore_order=True,
+    )
 
 
 def test_serialize_node__node_output(serialize_node):

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/ee/vellum_ee/workflows/display/utils/expressions.py

@@ -262,8 +262,27 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
                 "items": cast(JsonArray, serialized_items),  # list[JsonObject] -> JsonArray
             }
 
-    if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
-        raise ValueError("Nested references are not supported.")
+    if isinstance(value, dict):
+        serialized_entries = [
+            {"key": key, "value": serialize_value(display_context, val)} for key, val in value.items()
+        ]
+
+        # Check if all entries have constant values
+        if all(entry["value"]["type"] == "CONSTANT_VALUE" for entry in serialized_entries):
+            constant_entries = {}
+            for entry in serialized_entries:
+                entry_value = entry["value"]["value"]
+                constant_entries[entry["key"]] = entry_value["value"]
+
+            return {
+                "type": "CONSTANT_VALUE",
+                "value": {
+                    "type": "JSON",
+                    "value": constant_entries,
+                },
+            }
+        else:
+            return {"type": "DICTIONARY_REFERENCE", "entries": cast(JsonArray, serialized_entries)}
 
     if not isinstance(value, BaseDescriptor):
         vellum_value = primitive_to_vellum_value(value)

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/ee/vellum_ee/workflows/display/utils/vellum.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
+from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.client.types.vellum_value import VellumValue
@@ -117,13 +118,18 @@ def create_node_input_value_pointer_rule(
         workflow_input_display = display_context.global_workflow_input_displays[value]
         return InputVariablePointer(data=InputVariableData(input_variable_id=str(workflow_input_display.id)))
     if isinstance(value, VellumSecretReference):
-        workspace_secret = display_context.client.workspace_secrets.retrieve(
-            id=value.name,
-        )
+        try:
+            workspace_secret = display_context.client.workspace_secrets.retrieve(
+                id=value.name,
+            )
+            workspace_secret_id: Optional[str] = str(workspace_secret.id)
+        except ApiError:
+            workspace_secret_id = None
+
         return WorkspaceSecretPointer(
             data=WorkspaceSecretData(
                 type="STRING",
-                workspace_secret_id=str(workspace_secret.id),
+                workspace_secret_id=workspace_secret_id,
             ),
         )
     if isinstance(value, ExecutionCountReference):

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/ee/vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -587,3 +587,187 @@ def test_serialize_workflow__array_reference():
             },
         ],
     }
+
+
+def test_serialize_workflow__dict_values():
+    # GIVEN a node with a dictionary value
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            dict_value = {"key1": "value1", "key2": "value2"}
+            nested_dict_value = {
+                "key1": {"nested_key1": "value1", "nested_key2": "value2"},
+                "key2": {"nested_key1": "value1", "nested_key2": "value2"},
+            }
+            mixed_dict_value = {"key1": "value1", "key2": {"key3": "value3", "key4": "value4"}}
+
+    # AND a workflow that uses these outputs
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            dict_output = MyNode.Outputs.dict_value
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should serialize as a CONSTANT_VALUE
+    assert isinstance(data["workflow_raw_data"], dict)
+    assert isinstance(data["workflow_raw_data"]["nodes"], list)
+    my_node = next(
+        node for node in data["workflow_raw_data"]["nodes"] if isinstance(node, dict) and node["type"] == "GENERIC"
+    )
+
+    assert isinstance(my_node["outputs"], list)
+    outputs = my_node["outputs"]
+
+    dict_output = next(val for val in outputs if isinstance(val, dict) and val["name"] == "dict_value")
+    assert isinstance(dict_output, dict)
+    assert "value" in dict_output
+    assert dict_output["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "value": {"key1": "value1", "key2": "value2"}},
+    }
+
+    nested_dict_output = next(val for val in outputs if isinstance(val, dict) and val["name"] == "nested_dict_value")
+    assert isinstance(nested_dict_output, dict)
+    assert "value" in nested_dict_output
+    assert nested_dict_output["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {
+            "type": "JSON",
+            "value": {
+                "key1": {"nested_key1": "value1", "nested_key2": "value2"},
+                "key2": {"nested_key1": "value1", "nested_key2": "value2"},
+            },
+        },
+    }
+
+    mixed_dict_output = next(val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_dict_value")
+    assert isinstance(mixed_dict_output, dict)
+    assert "value" in mixed_dict_output
+    assert mixed_dict_output["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "value": {"key1": "value1", "key2": {"key3": "value3", "key4": "value4"}}},
+    }
+
+
+def test_serialize_workflow__dict_reference():
+    # GIVEN a node with a dictionary containing non-constant values (node references)
+    class FirstNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            value1: str
+
+    class SecondNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            # Dictionary containing a mix of constants and node references
+            mixed_dict = {
+                "key1": "constant1",
+                "key2": FirstNode.Outputs.value1,
+                "key3": "constant2",
+                "key4": FirstNode.Outputs.value1,
+            }
+            mixed_nested_dict = {
+                "key1": {"key1": "constant1", "key2": FirstNode.Outputs.value1},
+                "key2": {"key1": "constant2", "key2": FirstNode.Outputs.value1},
+            }
+
+    # AND a workflow that uses these outputs
+    class MyWorkflow(BaseWorkflow):
+        graph = FirstNode >> SecondNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            mixed_dict_output = SecondNode.Outputs.mixed_dict
+            mixed_nested_dict_output = SecondNode.Outputs.mixed_nested_dict
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should serialize as a CONSTANT_VALUE
+    assert isinstance(data["workflow_raw_data"], dict)
+    assert isinstance(data["workflow_raw_data"]["nodes"], list)
+    second_node = data["workflow_raw_data"]["nodes"][2]
+
+    assert isinstance(second_node, dict)
+    assert "outputs" in second_node
+    assert isinstance(second_node["outputs"], list)
+
+    outputs = second_node["outputs"]
+    mixed_dict_output = next(val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_dict")
+    assert isinstance(mixed_dict_output, dict)
+    assert "value" in mixed_dict_output
+    assert mixed_dict_output["value"] == {
+        "type": "DICTIONARY_REFERENCE",
+        "entries": [
+            {"key": "key1", "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant1"}}},
+            {
+                "key": "key2",
+                "value": {
+                    "type": "NODE_OUTPUT",
+                    "node_id": "13b4f5c0-e6aa-4ef9-9a1a-79476bc32500",
+                    "node_output_id": "50a6bc11-afb3-49f2-879c-b28f5e16d974",
+                },
+            },
+            {"key": "key3", "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant2"}}},
+            {
+                "key": "key4",
+                "value": {
+                    "type": "NODE_OUTPUT",
+                    "node_id": "13b4f5c0-e6aa-4ef9-9a1a-79476bc32500",
+                    "node_output_id": "50a6bc11-afb3-49f2-879c-b28f5e16d974",
+                },
+            },
+        ],
+    }
+
+    mixed_nested_dict_output = next(
+        val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_nested_dict"
+    )
+    assert isinstance(mixed_nested_dict_output, dict)
+    assert "value" in mixed_nested_dict_output
+    assert mixed_nested_dict_output["value"] == {
+        "type": "DICTIONARY_REFERENCE",
+        "entries": [
+            {
+                "key": "key1",
+                "value": {
+                    "type": "DICTIONARY_REFERENCE",
+                    "entries": [
+                        {
+                            "key": "key1",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant1"}},
+                        },
+                        {
+                            "key": "key2",
+                            "value": {
+                                "type": "NODE_OUTPUT",
+                                "node_id": "13b4f5c0-e6aa-4ef9-9a1a-79476bc32500",
+                                "node_output_id": "50a6bc11-afb3-49f2-879c-b28f5e16d974",
+                            },
+                        },
+                    ],
+                },
+            },
+            {
+                "key": "key2",
+                "value": {
+                    "type": "DICTIONARY_REFERENCE",
+                    "entries": [
+                        {
+                            "key": "key1",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant2"}},
+                        },
+                        {
+                            "key": "key2",
+                            "value": {
+                                "type": "NODE_OUTPUT",
+                                "node_id": "13b4f5c0-e6aa-4ef9-9a1a-79476bc32500",
+                                "node_output_id": "50a6bc11-afb3-49f2-879c-b28f5e16d974",
+                            },
+                        },
+                    ],
+                },
+            },
+        ],
+    }

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "0.14.48"
+version = "0.14.49"
 description = ""
 readme = "README.md"
 authors = []
@@ -93,7 +93,7 @@ plugins = [
     "vellum.plugins.vellum_mypy",
 ]
 check_untyped_defs = true
-exclude = "venv"
+exclude = ["venv", "examples/workflows"]
 overrides = [
     { module = "deepdiff.*", ignore_missing_imports = true },
     { module = "docker.*", ignore_missing_imports = true },
@@ -111,7 +111,7 @@ build-backend = "poetry.core.masonry.api"
 vellum = "vellum_cli:main"
 
 [tool.isort]
-skip = ["src/vellum/client", "tests/client", ".venv", "src/vellum/__init__.py"]
+skip = ["src/vellum/client", "tests/client", ".venv", "src/vellum/__init__.py", "examples/workflows"]
 multi_line_output = 3
 force_grid_wrap = 0
 line_length = 120
@@ -133,9 +133,9 @@ target-version = [
     "py38",
 ]
 include = ".*\\.py$"
-exclude = "(.git|venv|src/vellum/client|tests/client|scripts)"
+exclude = "(.git|venv|src/vellum/client|tests/client|scripts|examples/workflows)"
 
 [tool.flake8]
 max-line-length = 120
-exclude = ".git,.venv,ee/codegen/node_modules,ee/codegen_integration/fixtures,scripts,src/vellum/__init__.py,src/vellum/client,src/vellum/core,src/vellum/environment.py,src/vellum/errors,src/vellum/resources,src/vellum/types,tests/client"
+exclude = ".git,.venv,ee/codegen/node_modules,ee/codegen_integration/fixtures,scripts,src/vellum/__init__.py,src/vellum/client,src/vellum/core,src/vellum/environment.py,src/vellum/errors,src/vellum/resources,src/vellum/types,tests/client,examples/workflows"
 ignore = "E203, W503, E704"

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/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.14.48",
+            "X-Fern-SDK-Version": "0.14.49",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/src/vellum/workflows/nodes/displayable/code_execution_node/node.py

@@ -1,4 +1,5 @@
 import inspect
+import os
 from typing import Any, ClassVar, Dict, Generic, List, Optional, Sequence, Tuple, Type, TypeVar, cast, get_args
 
 from vellum import (
@@ -95,9 +96,9 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
     def run(self) -> Outputs:
         output_type = self.__class__.get_output_type()
-        code = self._resolve_code()
+        code, filepath = self._resolve_code()
         if not self.packages and self.runtime == "PYTHON_3_11_6":
-            logs, result = run_code_inline(code, self.code_inputs, output_type)
+            logs, result = run_code_inline(code, self.code_inputs, output_type, filepath)
             return self.Outputs(result=result, log=logs)
 
         else:
@@ -210,7 +211,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
         return compiled_inputs
 
-    def _resolve_code(self) -> str:
+    def _resolve_code(self) -> Tuple[str, str]:
         if self.code and self.filepath:
             raise NodeException(
                 message="Cannot specify both `code` and `filepath` for a CodeExecutionNode",
@@ -218,7 +219,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
             )
 
         if self.code:
-            return self.code
+            return self.code, f"{self.__class__.__name__}.code.py"
 
         if not self.filepath:
             raise NodeException(
@@ -235,4 +236,4 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                 code=WorkflowErrorCode.INVALID_INPUTS,
             )
 
-        return code
+        return code, os.path.join(os.path.dirname(root), self.filepath)

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/src/vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -850,7 +850,16 @@ def main(arg1: list) -> str:
         node.run()
 
     # AND the result should be the correct output
-    assert exc_info.value.message == "dict has no key: 'invalid'"
+    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'
+"""
+    )
 
 
 def test_run_node__execute_code__value_key_access():
@@ -933,3 +942,105 @@ def main(output: list[str]) -> list[str]:
 
     # 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": ""}
+
+
+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]
+"""
+        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
+"""
+    )

{vellum_ai-0.14.48 → vellum_ai-0.14.49}/src/vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -1,5 +1,7 @@
 import io
 import os
+import sys
+import traceback
 from typing import Any, Optional, Tuple, Union
 
 from pydantic import BaseModel
@@ -40,6 +42,7 @@ def run_code_inline(
     code: str,
     inputs: EntityInputsInterface,
     output_type: Any,
+    filepath: str,
 ) -> Tuple[str, Any]:
     log_buffer = io.StringIO()
 
@@ -75,11 +78,29 @@ def run_code_inline(
 __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()