vellum-ai 1.7.0__py3-none-any.whl → 1.7.2__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -27,10 +27,10 @@ class BaseClientWrapper:
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
-            "User-Agent": "vellum-ai/1.7.0",
+            "User-Agent": "vellum-ai/1.7.2",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "1.7.0",
+            "X-Fern-SDK-Version": "1.7.2",
             **(self.get_custom_headers() or {}),
         }
         if self._api_version is not None:

vellum/workflows/integrations/tests/test_vellum_integration_service.py

@@ -12,6 +12,7 @@ from vellum.workflows.types.definition import VellumIntegrationToolDetails
 
 def test_vellum_integration_service_get_tool_definition_success(vellum_client):
     """Test that tool definitions are successfully retrieved from Vellum API"""
+    # GIVEN a mock client configured to return a tool definition
     mock_client = vellum_client
     tool_definition_response = ComponentsSchemasComposioToolDefinition(
         integration=ToolDefinitionIntegration(
@@ -33,7 +34,6 @@ def test_vellum_integration_service_get_tool_definition_success(vellum_client):
         },
         output_parameters={},
     )
-
     mock_client.integrations.retrieve_integration_tool_definition.return_value = tool_definition_response
 
     # WHEN we request a tool definition
@@ -66,6 +66,7 @@ def test_vellum_integration_service_get_tool_definition_success(vellum_client):
 
 def test_vellum_integration_service_get_tool_definition_api_error(vellum_client):
     """Test that API errors are properly handled when retrieving tool definitions"""
+    # GIVEN a mock client configured to raise an exception
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
     mock_client.integrations.retrieve_integration_tool_definition.side_effect = Exception("Tool not found")
@@ -87,6 +88,7 @@ def test_vellum_integration_service_get_tool_definition_api_error(vellum_client)
 
 def test_vellum_integration_service_execute_tool_success(vellum_client):
     """Test that tools are successfully executed via Vellum API"""
+    # GIVEN a mock client configured to return a successful response
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
 
@@ -96,7 +98,6 @@ def test_vellum_integration_service_execute_tool_success(vellum_client):
         "issue_id": 123,
         "issue_url": "https://github.com/user/repo/issues/123",
     }
-
     mock_client.integrations.execute_integration_tool.return_value = mock_response
 
     # WHEN we execute a tool with valid arguments
@@ -132,6 +133,7 @@ def test_vellum_integration_service_execute_tool_success(vellum_client):
 
 def test_vellum_integration_service_execute_tool_api_error(vellum_client):
     """Test that execution errors are properly handled"""
+    # GIVEN a mock client configured to raise an exception
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
     mock_client.integrations.execute_integration_tool.side_effect = Exception("Authentication failed")
@@ -154,12 +156,12 @@ def test_vellum_integration_service_execute_tool_api_error(vellum_client):
 
 def test_vellum_integration_service_execute_tool_empty_response(vellum_client):
     """Test that empty response data is handled gracefully"""
+    # GIVEN a mock client configured to return an empty response
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
 
     mock_response = mock.MagicMock()
     mock_response.data = {}
-
     mock_client.integrations.execute_integration_tool.return_value = mock_response
 
     # WHEN we execute a tool that returns empty data
@@ -180,6 +182,7 @@ def test_vellum_integration_service_execute_tool_empty_response(vellum_client):
 
 def test_vellum_integration_service_multiple_tool_executions(vellum_client):
     """Test that the service handles multiple sequential tool executions"""
+    # GIVEN a mock client configured to return different responses for multiple calls
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
 
@@ -214,16 +217,18 @@ def test_vellum_integration_service_multiple_tool_executions(vellum_client):
     assert mock_client.integrations.execute_integration_tool.call_count == 2
 
 
-def test_vellum_integration_service_execute_tool_structured_403_error(vellum_client):
-    """Test that structured 403 responses with integration details are properly parsed"""
+def test_vellum_integration_service_execute_tool_structured_403_with_integration(vellum_client):
+    """Test structured 403 responses with integration field (current backend format)"""
     from vellum.client.core.api_error import ApiError
     from vellum.workflows.errors.types import WorkflowErrorCode
 
+    # GIVEN a mock client configured to raise a structured 403 error with integration
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
 
-    # Mock structured 403 response matching PR #14857 format
+    # Mock current backend structure with integration as direct field
     structured_error_body = {
+        "code": "INTEGRATION_CREDENTIALS_UNAVAILABLE",
         "message": "You must authenticate with this integration before you can execute this tool.",
         "integration": {
             "id": "550e8400-e29b-41d4-a716-446655440000",
@@ -231,15 +236,13 @@ def test_vellum_integration_service_execute_tool_structured_403_error(vellum_cli
             "name": "GITHUB",
         },
     }
-
     mock_client.integrations.execute_integration_tool.side_effect = ApiError(
         status_code=403,
         body=structured_error_body,
     )
 
-    service = VellumIntegrationService(client=mock_client)
-
     # WHEN we attempt to execute a tool without credentials
+    service = VellumIntegrationService(client=mock_client)
     with pytest.raises(NodeException) as exc_info:
         service.execute_tool(
             integration="GITHUB",
@@ -266,19 +269,18 @@ def test_vellum_integration_service_execute_tool_legacy_403_error(vellum_client)
     from vellum.client.core.api_error import ApiError
     from vellum.workflows.errors.types import WorkflowErrorCode
 
+    # GIVEN a mock client configured to raise a legacy 403 error
     mock_client = vellum_client
     mock_client.integrations = mock.MagicMock()
 
-    # Mock legacy 403 response format (just detail field)
     legacy_error_body = {"detail": "You do not have permission to execute this tool."}
-
     mock_client.integrations.execute_integration_tool.side_effect = ApiError(
         status_code=403,
         body=legacy_error_body,
     )
 
+    # WHEN we attempt to execute a tool that returns a legacy 403 error
     service = VellumIntegrationService(client=mock_client)
-
     with pytest.raises(NodeException) as exc_info:
         service.execute_tool(
             integration="GITHUB",

vellum/workflows/integrations/vellum_integration_service.py

@@ -95,15 +95,15 @@ class VellumIntegrationService:
         except ApiError as e:
             # Handle structured 403 credential error responses
             if e.status_code == 403 and isinstance(e.body, dict):
-                if "integration" in e.body and "message" in e.body:
-                    integration_details = e.body["integration"]
-                    error_message = e.body["message"]
-
-                    # Keep integration details nested under "integration" key to keep raw_data raw
-                    # and allow for future expansion
-                    raw_data = {
-                        "integration": integration_details,
-                    }
+                # Check for backend structure with integration as direct field
+                integration_from_backend = e.body.get("integration")
+                if integration_from_backend:
+                    error_message = e.body.get(
+                        "message", "You must authenticate with this integration before you can execute this tool."
+                    )
+
+                    # Wrap integration in raw_data for frontend consumption
+                    raw_data = {"integration": integration_from_backend}
 
                     raise NodeException(
                         message=error_message,

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

@@ -117,13 +117,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
@@ -134,6 +128,24 @@ 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():

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

@@ -6,7 +6,11 @@ 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
@@ -792,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]):

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

@@ -22,8 +22,8 @@ from vellum.workflows.nodes.displayable.tool_calling_node.utils import (
 )
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
 from vellum.workflows.state.context import WorkflowContext
-from vellum.workflows.types.core import EntityInputsInterface, Tool
-from vellum.workflows.types.definition import MCPServer
+from vellum.workflows.types.core import EntityInputsInterface
+from vellum.workflows.types.definition import MCPServer, Tool
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.utils.functions import compile_mcp_tool_definition, get_mcp_tool_name
 from vellum.workflows.workflows.event_filters import all_workflow_event_filter

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

@@ -31,12 +31,14 @@ from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.state import BaseState
 from vellum.workflows.state.encoder import DefaultStateEncoder
-from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior, Tool, ToolBase
+from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.definition import (
     ComposioToolDefinition,
     DeploymentDefinition,
     MCPServer,
     MCPToolDefinition,
+    Tool,
+    ToolBase,
     VellumIntegrationToolDefinition,
 )
 from vellum.workflows.types.generics import is_workflow_class
@@ -203,7 +205,7 @@ class FunctionNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
             raise NodeException(
                 message=f"Error executing function '{function_name}': {str(e)}",
                 code=WorkflowErrorCode.NODE_EXECUTION,
-            )
+            ) from e
 
         # Add the result to the chat history
         self._add_function_result_to_chat_history(result, self.state)
@@ -233,7 +235,7 @@ class ComposioNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
             raise NodeException(
                 message=f"Error executing Composio tool '{self.composio_tool.action}': {str(e)}",
                 code=WorkflowErrorCode.NODE_EXECUTION,
-            )
+            ) from e
 
         # Add result to chat history
         self._add_function_result_to_chat_history(result, self.state)
@@ -256,7 +258,7 @@ class MCPNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
             raise NodeException(
                 message=f"Error executing MCP tool '{self.mcp_tool.name}': {str(e)}",
                 code=WorkflowErrorCode.NODE_EXECUTION,
-            )
+            ) from e
 
         # Add result to chat history
         self._add_function_result_to_chat_history(result, self.state)
@@ -281,11 +283,18 @@ class VellumIntegrationNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
                 tool_name=self.vellum_integration_tool.name,
                 arguments=arguments,
             )
+        except NodeException as e:
+            # Preserve original error code and raw_data while adding context
+            raise NodeException(
+                message=f"Error executing Vellum Integration tool '{self.vellum_integration_tool.name}': {e.message}",
+                code=e.code,
+                raw_data=e.raw_data,
+            ) from e
         except Exception as e:
             raise NodeException(
                 message=f"Error executing Vellum Integration tool '{self.vellum_integration_tool.name}': {str(e)}",
                 code=WorkflowErrorCode.NODE_EXECUTION,
-            )
+            ) from e
 
         # Add result to chat history
         self._add_function_result_to_chat_history(result, self.state)

vellum/workflows/resolvers/resolver.py

@@ -1,9 +1,7 @@
 import logging
 from uuid import UUID
-from typing import Iterator, List, Optional, Tuple, Type, Union
+from typing import Iterator, Optional, Type, Union
 
-from vellum.client.types.vellum_span import VellumSpan
-from vellum.client.types.workflow_execution_initiated_event import WorkflowExecutionInitiatedEvent
 from vellum.workflows.events.workflow import WorkflowEvent
 from vellum.workflows.nodes.utils import cast_to_output_type
 from vellum.workflows.resolvers.base import BaseWorkflowResolver
@@ -20,38 +18,6 @@ class VellumResolver(BaseWorkflowResolver):
     def get_state_snapshot_history(self) -> Iterator[BaseState]:
         return iter([])
 
-    def _find_previous_and_root_span(
-        self, execution_id: str, spans: List[VellumSpan]
-    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
-        previous_trace_id: Optional[str] = None
-        root_trace_id: Optional[str] = None
-        previous_span_id: Optional[str] = None
-        root_span_id: Optional[str] = None
-
-        for span in spans:
-            # Look for workflow execution spans with matching ID first
-            if span.name == "workflow.execution" and span.span_id == execution_id:
-                # Find the WorkflowExecutionInitiatedEvent in the span's events
-                initiated_event = next(
-                    (event for event in span.events if isinstance(event, WorkflowExecutionInitiatedEvent)), None
-                )
-                if initiated_event:
-                    previous_trace_id = initiated_event.trace_id
-                    previous_span_id = initiated_event.span_id
-                    links = initiated_event.links
-                    if links:
-                        root_span = next((link for link in links if link.type == "ROOT_SPAN"), None)
-                        if root_span:
-                            root_trace_id = root_span.trace_id
-                            root_span_id = root_span.span_context.span_id
-                    else:
-                        # no links means this is the first execution
-                        root_trace_id = initiated_event.trace_id
-                        root_span_id = initiated_event.span_id
-                    break
-
-        return previous_trace_id, root_trace_id, previous_span_id, root_span_id
-
     def _deserialize_state(self, state_data: dict, state_class: Type[BaseState]) -> BaseState:
         """Deserialize state data with proper type conversion for complex types like List[ChatMessage]."""
         converted_data = {}
@@ -79,18 +45,19 @@ class VellumResolver(BaseWorkflowResolver):
             return None
 
         client = self._context.vellum_client
-        response = client.workflow_executions.retrieve_workflow_execution_detail(
-            execution_id=previous_execution_id,
+        response = client.workflows.retrieve_state(
+            span_id=previous_execution_id,
         )
 
         if response.state is None:
             return None
 
-        previous_trace_id, root_trace_id, previous_span_id, root_span_id = self._find_previous_and_root_span(
-            previous_execution_id, response.spans
-        )
-
-        if previous_trace_id is None or root_trace_id is None or previous_span_id is None or root_span_id is None:
+        if (
+            response.previous_trace_id is None
+            or response.root_trace_id is None
+            or response.previous_span_id is None
+            or response.root_span_id is None
+        ):
             logger.warning("Could not find required execution events for state loading")
             return None
 
@@ -106,8 +73,8 @@ class VellumResolver(BaseWorkflowResolver):
 
         return LoadStateResult(
             state=state,
-            previous_trace_id=previous_trace_id,
-            previous_span_id=previous_span_id,
-            root_trace_id=root_trace_id,
-            root_span_id=root_span_id,
+            previous_trace_id=response.previous_trace_id,
+            previous_span_id=response.previous_span_id,
+            root_trace_id=response.root_trace_id,
+            root_span_id=response.root_span_id,
         )

vellum/workflows/resolvers/tests/test_resolver.py

@@ -4,14 +4,7 @@ from uuid import uuid4
 from typing import List
 
 from vellum import ChatMessage
-from vellum.client.types.span_link import SpanLink
-from vellum.client.types.vellum_code_resource_definition import VellumCodeResourceDefinition
-from vellum.client.types.workflow_execution_detail import WorkflowExecutionDetail
-from vellum.client.types.workflow_execution_initiated_body import WorkflowExecutionInitiatedBody
-from vellum.client.types.workflow_execution_initiated_event import WorkflowExecutionInitiatedEvent
-from vellum.client.types.workflow_execution_span import WorkflowExecutionSpan
-from vellum.client.types.workflow_execution_span_attributes import WorkflowExecutionSpanAttributes
-from vellum.client.types.workflow_parent_context import WorkflowParentContext
+from vellum.client.types.workflow_resolved_state import WorkflowResolvedState
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.resolvers.resolver import VellumResolver
@@ -24,7 +17,6 @@ def test_load_state_with_context_success():
     """Test load_state successfully loads state when context and client are available."""
     resolver = VellumResolver()
     execution_id = uuid4()
-    root_execution_id = uuid4()
 
     class TestState(BaseState):
         test_key: str = "test_value"
@@ -50,60 +42,24 @@ def test_load_state_with_context_success():
         },
     }
 
-    mock_workflow_definition = VellumCodeResourceDefinition(
-        name="TestWorkflow", module=["test", "module"], id=str(uuid4())
-    )
-
-    mock_body = WorkflowExecutionInitiatedBody(workflow_definition=mock_workflow_definition, inputs={})
-
     previous_trace_id = str(uuid4())
+    previous_span_id = str(uuid4())
     root_trace_id = str(uuid4())
+    root_span_id = str(uuid4())
 
-    previous_invocation = WorkflowExecutionInitiatedEvent(
-        id=str(uuid4()),
+    mock_response = WorkflowResolvedState(
+        trace_id=str(uuid4()),
         timestamp=datetime.now(),
-        trace_id=previous_trace_id,
         span_id=str(execution_id),
-        body=mock_body,
-        links=[
-            SpanLink(
-                trace_id=previous_trace_id,
-                type="PREVIOUS_SPAN",
-                span_context=WorkflowParentContext(workflow_definition=mock_workflow_definition, span_id=str(uuid4())),
-            ),
-            SpanLink(
-                trace_id=root_trace_id,
-                type="ROOT_SPAN",
-                span_context=WorkflowParentContext(
-                    workflow_definition=mock_workflow_definition, span_id=str(root_execution_id)
-                ),
-            ),
-        ],
-    )
-
-    root_invocation = WorkflowExecutionInitiatedEvent(
-        id=str(uuid4()),
-        timestamp=datetime.now(),
-        trace_id=root_trace_id,
-        span_id=str(root_execution_id),
-        body=mock_body,
-        links=None,  # Root invocation has no links
-    )
-
-    mock_span = WorkflowExecutionSpan(
-        span_id=str(execution_id),  # Use the actual execution_id
-        start_ts=datetime.now(),
-        end_ts=datetime.now(),
-        attributes=WorkflowExecutionSpanAttributes(label="Test Workflow", workflow_id=str(uuid4())),
-        events=[previous_invocation, root_invocation],
-    )
-
-    mock_response = WorkflowExecutionDetail(
-        span_id="test-span-id", start=datetime.now(), inputs=[], outputs=[], spans=[mock_span], state=state_dict
+        state=state_dict,
+        previous_trace_id=previous_trace_id,
+        previous_span_id=previous_span_id,
+        root_trace_id=root_trace_id,
+        root_span_id=root_span_id,
     )
 
     mock_client = Mock()
-    mock_client.workflow_executions.retrieve_workflow_execution_detail.return_value = mock_response
+    mock_client.workflows.retrieve_state.return_value = mock_response
 
     # AND context with the test workflow class is set up
     context = WorkflowContext(vellum_client=mock_client)
@@ -123,21 +79,18 @@ def test_load_state_with_context_success():
     assert str(result.state.meta.span_id) != prev_span_id
 
     # AND should have span link info
-    assert result.previous_trace_id == previous_invocation.trace_id
-    assert result.previous_span_id == previous_invocation.span_id
-    assert result.root_trace_id == root_invocation.trace_id
-    assert result.root_span_id == root_invocation.span_id
+    assert result.previous_trace_id == previous_trace_id
+    assert result.previous_span_id == previous_span_id
+    assert result.root_trace_id == root_trace_id
+    assert result.root_span_id == root_span_id
 
-    mock_client.workflow_executions.retrieve_workflow_execution_detail.assert_called_once_with(
-        execution_id=str(execution_id)
-    )
+    mock_client.workflows.retrieve_state.assert_called_once_with(span_id=str(execution_id))
 
 
 def test_load_state_with_chat_message_list():
     """Test load_state successfully loads state with chat_history containing ChatMessage list."""
     resolver = VellumResolver()
     execution_id = uuid4()
-    root_execution_id = uuid4()
 
     class TestStateWithChatHistory(BaseState):
         test_key: str = "test_value"
@@ -169,60 +122,24 @@ def test_load_state_with_chat_message_list():
         },
     }
 
-    mock_workflow_definition = VellumCodeResourceDefinition(
-        name="TestWorkflow", module=["test", "module"], id=str(uuid4())
-    )
-
-    mock_body = WorkflowExecutionInitiatedBody(workflow_definition=mock_workflow_definition, inputs={})
-
     previous_trace_id = str(uuid4())
+    previous_span_id = str(uuid4())
     root_trace_id = str(uuid4())
+    root_span_id = str(uuid4())
 
-    previous_invocation = WorkflowExecutionInitiatedEvent(
-        id=str(uuid4()),
-        timestamp=datetime.now(),
-        trace_id=previous_trace_id,
-        span_id=str(execution_id),
-        body=mock_body,
-        links=[
-            SpanLink(
-                trace_id=previous_trace_id,
-                type="PREVIOUS_SPAN",
-                span_context=WorkflowParentContext(workflow_definition=mock_workflow_definition, span_id=str(uuid4())),
-            ),
-            SpanLink(
-                trace_id=root_trace_id,
-                type="ROOT_SPAN",
-                span_context=WorkflowParentContext(
-                    workflow_definition=mock_workflow_definition, span_id=str(root_execution_id)
-                ),
-            ),
-        ],
-    )
-
-    root_invocation = WorkflowExecutionInitiatedEvent(
-        id=str(uuid4()),
+    mock_response = WorkflowResolvedState(
+        trace_id=str(uuid4()),
         timestamp=datetime.now(),
-        trace_id=root_trace_id,
-        span_id=str(root_execution_id),
-        body=mock_body,
-        links=None,
-    )
-
-    mock_span = WorkflowExecutionSpan(
         span_id=str(execution_id),
-        start_ts=datetime.now(),
-        end_ts=datetime.now(),
-        attributes=WorkflowExecutionSpanAttributes(label="Test Workflow", workflow_id=str(uuid4())),
-        events=[previous_invocation, root_invocation],
-    )
-
-    mock_response = WorkflowExecutionDetail(
-        span_id="test-span-id", start=datetime.now(), inputs=[], outputs=[], spans=[mock_span], state=state_dict
+        state=state_dict,
+        previous_trace_id=previous_trace_id,
+        previous_span_id=previous_span_id,
+        root_trace_id=root_trace_id,
+        root_span_id=root_span_id,
     )
 
     mock_client = Mock()
-    mock_client.workflow_executions.retrieve_workflow_execution_detail.return_value = mock_response
+    mock_client.workflows.retrieve_state.return_value = mock_response
 
     # AND context with the test workflow class is set up
     context = WorkflowContext(vellum_client=mock_client)
@@ -247,6 +164,4 @@ def test_load_state_with_chat_message_list():
     assert result.state.chat_history[2].role == "USER"
     assert result.state.chat_history[2].text == "What can you help me with?"
 
-    mock_client.workflow_executions.retrieve_workflow_execution_detail.assert_called_once_with(
-        execution_id=str(execution_id)
-    )
+    mock_client.workflows.retrieve_state.assert_called_once_with(span_id=str(execution_id))

vellum/workflows/state/base.py

@@ -442,14 +442,24 @@ class StateMeta(UniversalBaseModel):
         if not memo:
             memo = {}
 
+        node_output_keys = list(self.node_outputs.keys())
         new_node_outputs = {
-            descriptor: value if isinstance(value, Queue) else deepcopy(value, memo)
-            for descriptor, value in self.node_outputs.items()
+            descriptor: (
+                self.node_outputs[descriptor]
+                if isinstance(self.node_outputs[descriptor], Queue)
+                else deepcopy(self.node_outputs[descriptor], memo)
+            )
+            for descriptor in node_output_keys
         }
 
+        external_input_keys = list(self.external_inputs.keys())
         new_external_inputs = {
-            descriptor: value if isinstance(value, Queue) else deepcopy(value, memo)
-            for descriptor, value in self.external_inputs.items()
+            descriptor: (
+                self.external_inputs[descriptor]
+                if isinstance(self.external_inputs[descriptor], Queue)
+                else deepcopy(self.external_inputs[descriptor], memo)
+            )
+            for descriptor in external_input_keys
         }
 
         memo[id(self.node_outputs)] = new_node_outputs

vellum/workflows/types/core.py

@@ -1,11 +1,8 @@
 from enum import Enum
 from typing import (  # type: ignore[attr-defined]
-    TYPE_CHECKING,
     Any,
-    Callable,
     Dict,
     List,
-    Type,
     Union,
     _GenericAlias,
     _SpecialGenericAlias,
@@ -13,16 +10,6 @@ from typing import (  # type: ignore[attr-defined]
 )
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
-from vellum.workflows.types.definition import (
-    ComposioToolDefinition,
-    DeploymentDefinition,
-    MCPServer,
-    VellumIntegrationToolDefinition,
-)
-
-if TYPE_CHECKING:
-    from vellum.workflows.workflows.base import BaseWorkflow
-
 
 JsonArray = List["Json"]
 JsonObject = Dict[str, "Json"]
@@ -53,14 +40,3 @@ class ConditionType(Enum):
     IF = "IF"
     ELIF = "ELIF"
     ELSE = "ELSE"
-
-
-# Type alias for functions that can be called in tool calling nodes
-ToolBase = Union[
-    Callable[..., Any],
-    DeploymentDefinition,
-    Type["BaseWorkflow"],
-    ComposioToolDefinition,
-    VellumIntegrationToolDefinition,
-]
-Tool = Union[ToolBase, MCPServer]

vellum/workflows/types/definition.py

@@ -2,7 +2,7 @@ import importlib
 import inspect
 from types import FrameType
 from uuid import UUID
-from typing import Annotated, Any, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Annotated, Any, Callable, Dict, List, Literal, Optional, Type, Union
 
 from pydantic import BeforeValidator, SerializationInfo, model_serializer
 
@@ -13,6 +13,9 @@ from vellum.client.types.vellum_variable import VellumVariable
 from vellum.workflows.constants import AuthorizationType, VellumIntegrationProviderType
 from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
 
 def serialize_type_encoder(obj: type) -> Dict[str, Any]:
     return {
@@ -216,3 +219,14 @@ class MCPToolDefinition(UniversalBaseModel):
     server: MCPServer
     description: Optional[str] = None
     parameters: Dict[str, Any] = {}
+
+
+# Type alias for functions that can be called in tool calling nodes
+ToolBase = Union[
+    Callable[..., Any],
+    DeploymentDefinition,
+    Type["BaseWorkflow"],
+    ComposioToolDefinition,
+    VellumIntegrationToolDefinition,
+]
+Tool = Union[ToolBase, MCPServer]

{vellum_ai-1.7.0.dist-info → vellum_ai-1.7.2.dist-info}/METADATA

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

{vellum_ai-1.7.0.dist-info → vellum_ai-1.7.2.dist-info}/RECORD

@@ -22,7 +22,7 @@ vellum_cli/tests/test_ping.py,sha256=b3aQLd-N59_8w2rRiWqwpB1rlHaKEYVbAj1Y3hi7A-g
 vellum_cli/tests/test_pull.py,sha256=e2XHzcHIx9k-FyuNAl7wMSNsSSebPGyP6U05JGcddFs,49447
 vellum_cli/tests/test_push.py,sha256=2MjkNKr_9Guv5Exjsm3L1BeVXmPkKUcCSiKnp90HgW4,41996
 vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/assets/node-definitions.json,sha256=tinXep1k-FicJWSPQb5mlWRkdKul1WJtG4o94SqD4ds,29800
+vellum_ee/assets/node-definitions.json,sha256=UkHixt8WAnHmupooOrtpxYCaNG-5pRon1c8QUlc0Vhk,30754
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/scripts/generate_node_definitions.py,sha256=FOYQsXIqU45I0OAcsyZUGODF9JK44yunf58rR6YaAdA,3303
 vellum_ee/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -55,7 +55,7 @@ vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py,sha256=Ue727X
 vellum_ee/workflows/display/nodes/vellum/retry_node.py,sha256=5xv5OR4xRuI1R5yqJDZTNdCFYY-z8PkTdpWM4ziGjw0,3192
 vellum_ee/workflows/display/nodes/vellum/search_node.py,sha256=ZHqkwzCRsM0TkZmh3ePHreZAoQXeT-SFS7zbUrSjMsw,12037
 vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=hC5hTSRf1f9ppiZ_wXzpcp-jl0fwNAUB9PY8s0AykHs,3131
-vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=QExp4l2KMNBdqMVDL4CG_iFMpFf3RYiTP2q5P2DmEFk,3178
+vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=sWupmR7Ri2ZeCLx8713JbuawabR_kpWp_-HwT1tCrtU,3163
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/vellum/tests/test_api_node.py,sha256=DQAtsabvn6BE6xWwKNHzMOppzoy1-1dssNnrwbHUdRU,1490
 vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py,sha256=6_1yVYftVnqXJn3hLUGHlcHvbgKbQgBfS6vYQ6R79oc,10891
@@ -79,7 +79,7 @@ vellum_ee/workflows/display/tests/test_base_workflow_display.py,sha256=-LYQw5nWq
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=Y-ajeT65b5varmrZCw6L3hir4hJCFq-eO0jZfRcrs7g,1886
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=bD9AtBo7Pz1Drbh420NR0VSiHBfExw_UtNRPCiAqki0,13967
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=wdVaFvxicj48Kj-6VUlu61KR0oSArLTjownRi2p0NWQ,14941
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=bPTwkLpB7trFLpAaDvXMfMP0c9H1u_c1cdnj7K-gtnw,24962
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=SwXRzjdoEZLvkzaRMvRV8_UqbBm0EB_UtAHD_zXKZBY,6303
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=jubGYgLJImOqILd5LjnYJ4B1UMIrToDrQbPZOvaQCX4,40035
@@ -126,7 +126,7 @@ vellum_ee/workflows/display/utils/tests/test_events.py,sha256=42IEBnMbaQrH8gigw5
 vellum_ee/workflows/display/utils/vellum.py,sha256=Bt7kdLdXoBsHn5dVEY2uKcF542VL09jwu8J_30rl2vk,6413
 vellum_ee/workflows/display/vellum.py,sha256=J2mdJZ1sdLW535DDUkq_Vm8Z572vhuxHxVZF9deKSdk,391
 vellum_ee/workflows/display/workflows/__init__.py,sha256=JTB9ObEV3l4gGGdtfBHwVJtTTKC22uj-a-XjTVwXCyA,148
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=FAsEeqkQu4xGvdZ4Ri5H7LDEISQYziidPc2QHzbOVWs,44311
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=uP3pqSloGjmUYBvVfKP0skZXJnqZxm7TEVxSLdrzU5c,44682
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=gxz76AeCqgAZ9D2lZeTiZzxY9eMgn3qOSfVgiqYcOh8,2028
 vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=lg-c_3P3ldtqWq2VFsk_2Mkn3pVdXWuT59QpH7QwXVs,39764
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -160,7 +160,7 @@ vellum/client/README.md,sha256=flqu57ubZNTfpq60CdLtJC9gp4WEkyjb_n_eZ4OYf9w,6497
 vellum/client/__init__.py,sha256=rMnKRqL5-356SBc-rfm56MkO87PuAi2mtcfBszcJU1M,74316
 vellum/client/core/__init__.py,sha256=lTcqUPXcx4112yLDd70RAPeqq6tu3eFMe1pKOqkW9JQ,1562
 vellum/client/core/api_error.py,sha256=44vPoTyWN59gonCIZMdzw7M1uspygiLnr3GNFOoVL2Q,614
-vellum/client/core/client_wrapper.py,sha256=7jdvkc9bp52v55uLBImLe_fMXRZUgS2HqLAzuk-JDRg,2840
+vellum/client/core/client_wrapper.py,sha256=4Sfjvux7hceG2zs5IHHn9McKc2YE_WIb9ctx_HfCyzA,2840
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=d4NNbX8XvXP32z8KpK2Xovv33nFfruIrpz0QWxlgpZk,2663
 vellum/client/core/force_multipart.py,sha256=awxh5MtcRYe74ehY8U76jzv6fYM_w_D3Rur7KQQzSDk,429
@@ -1876,8 +1876,8 @@ vellum/workflows/integrations/composio_service.py,sha256=rSliaZtNiBcDSvDxz9k5i1K
 vellum/workflows/integrations/mcp_service.py,sha256=9DYb8dg2_kgc1UOu830kxhaFlt9yTbhKPhK3L6kb1t4,9831
 vellum/workflows/integrations/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/integrations/tests/test_mcp_service.py,sha256=q_DYrDkIqI4sQBNgID4YdbM4e9tneLVWY8YmI4R26d8,8859
-vellum/workflows/integrations/tests/test_vellum_integration_service.py,sha256=n_n42FnwbR7YMd1sEVMphdOUWsWXCPw_caY6FnxNxAU,11322
-vellum/workflows/integrations/vellum_integration_service.py,sha256=JFkrXBhkTEn_tXnRbjt289Z56_SLDSsuU7Msx2Nh9AM,5192
+vellum/workflows/integrations/tests/test_vellum_integration_service.py,sha256=QGJmaW5EF7E2fkZotd4rc2HItQc-1z3cpgwaKUFWpgg,11956
+vellum/workflows/integrations/vellum_integration_service.py,sha256=qhFoLzHlMli1PC8oh5phvWuSpJ9IqL1g2eaGhypBTqs,5266
 vellum/workflows/logging.py,sha256=_a217XogktV4Ncz6xKFz7WfYmZAzkfVRVuC0rWob8ls,437
 vellum/workflows/nodes/__init__.py,sha256=zymtc3_iW2rFmMR-sayTLuN6ZsAw8VnJweWPsjQk2-Q,1197
 vellum/workflows/nodes/bases/__init__.py,sha256=cniHuz_RXdJ4TQgD8CBzoiKDiPxg62ErdVpCbWICX64,58
@@ -1931,11 +1931,11 @@ vellum/workflows/nodes/displayable/bases/tests/test_utils.py,sha256=eqdqbKNRWVMD
 vellum/workflows/nodes/displayable/bases/types.py,sha256=C37B2Qh2YP7s7pUjd-EYKc2Zl1TbnCgI_mENuUSb8bo,1706
 vellum/workflows/nodes/displayable/bases/utils.py,sha256=X1YSPmbBlxDrTAw6dy2-9HrIG8Vb_J-2k1lP3i-SOsk,6951
 vellum/workflows/nodes/displayable/code_execution_node/__init__.py,sha256=0FLWMMktpzSnmBMizQglBpcPrP80fzVsoJwJgf822Cg,76
-vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=t4AuqyGOSgitYUUs6Tqzo7lDeyOyJ5g8x7Q9j4UkH7M,10426
+vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=JkJ_p9KKmrkUqc8TPcxfPoZ1AiEnWWwrdzLj-0TS9ds,10985
 vellum/workflows/nodes/displayable/code_execution_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/main.py,sha256=5QsbmkzSlSbcbWTG_JmIqcP-JNJzOPTKxGzdHos19W4,79
-vellum/workflows/nodes/displayable/code_execution_node/tests/test_node.py,sha256=RQy7SsA39kTJAudzpYNN3jK_kuq6bWNZUld54H5pRZA,40966
+vellum/workflows/nodes/displayable/code_execution_node/tests/test_node.py,sha256=eCWBaAPzneDiOQe1xFyyb9h9LeQQOAdjRY1RJS0z8sE,44500
 vellum/workflows/nodes/displayable/code_execution_node/utils.py,sha256=0F_4PVvB7vjsV0RS48Brv_4djebWOMR_zgzHxzA9iV4,3308
 vellum/workflows/nodes/displayable/conditional_node/__init__.py,sha256=AS_EIqFdU1F9t8aLmbZU-rLh9ry6LCJ0uj0D8F0L5Uw,72
 vellum/workflows/nodes/displayable/conditional_node/node.py,sha256=2g32hWosE3zwVaK2UPFnVEer1Nbn04s3friF2rGztmE,1195
@@ -1975,13 +1975,13 @@ vellum/workflows/nodes/displayable/tests/test_search_node_error_handling.py,sha2
 vellum/workflows/nodes/displayable/tests/test_search_node_wth_text_output.py,sha256=VepO5z1277c1y5N6LLIC31nnWD1aak2m5oPFplfJHHs,6935
 vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py,sha256=Bjv-wZyFgNaVZb9KEMMZd9lFoLzbPEPjEMpANizMZw4,2413
 vellum/workflows/nodes/displayable/tool_calling_node/__init__.py,sha256=3n0-ysmFKsr40CVxPthc0rfJgqVJeZuUEsCmYudLVRg,117
-vellum/workflows/nodes/displayable/tool_calling_node/node.py,sha256=RxwBjNMmpRBu2G6fnu37MvqRVxa7ZQzOPtqccCYlCqI,8420
+vellum/workflows/nodes/displayable/tool_calling_node/node.py,sha256=pwK9q1blgRv9Mz_LY-fv_hxHfOOkuKxJ6AgRlorUrXk,8420
 vellum/workflows/nodes/displayable/tool_calling_node/state.py,sha256=CcBVb_YtwfSSka4ze678k6-qwmzMSfjfVP8_Y95feSo,302
 vellum/workflows/nodes/displayable/tool_calling_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/tool_calling_node/tests/test_composio_service.py,sha256=in1fbEz5x1tx3uKv9YXdvOncsHucNL8Ro6Go7lBuuOQ,8962
 vellum/workflows/nodes/displayable/tool_calling_node/tests/test_node.py,sha256=GZoeybB9uM7ai8sBLAtUMHrMVgh-WrJDWrKZci6feDs,11892
 vellum/workflows/nodes/displayable/tool_calling_node/tests/test_utils.py,sha256=EmKFA-ELdTzlK0xMqWnuSZPoGNLYCwk6b0amTqirZo0,11305
-vellum/workflows/nodes/displayable/tool_calling_node/utils.py,sha256=7xNiNPD4vzd-4NuIaeCZhs0AshkuxlWVdx1NnbX-fdg,23722
+vellum/workflows/nodes/displayable/tool_calling_node/utils.py,sha256=NDp2rEU-1QACxrSplocDKlPzILczzvPygkUOrysjU7Y,24109
 vellum/workflows/nodes/displayable/web_search_node/__init__.py,sha256=8FOnEP-n-U68cvxTlJW9wphIAGHq5aqjzLM-DoSSXnU,61
 vellum/workflows/nodes/displayable/web_search_node/node.py,sha256=NQYux2bOtuBF5E4tn-fXi5y3btURPRrNqMSM9MAZYI4,5091
 vellum/workflows/nodes/displayable/web_search_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -2017,14 +2017,14 @@ vellum/workflows/references/vellum_secret.py,sha256=Od4d19a5yletWMqNfJR5d_mZQUkV
 vellum/workflows/references/workflow_input.py,sha256=W3rOK1EPd2gYHb04WJwmNm1CUSdvZ9LKrs8RMKxACBs,1751
 vellum/workflows/resolvers/__init__.py,sha256=eH6hTvZO4IciDaf_cf7aM2vs-DkBDyJPycOQevJxQnI,82
 vellum/workflows/resolvers/base.py,sha256=wrQiSC02Bw4-dBwgFjJIHsjpe-4xz4rUJs_1RdErKA0,1164
-vellum/workflows/resolvers/resolver.py,sha256=yK-oY0HDsFJcjlNKAm3vpsPKRIFerIh59FcTwuEN5RY,4839
-vellum/workflows/resolvers/tests/test_resolver.py,sha256=jXkJBb9SwtoH__bBN-ECohpyD0aTIB9ErEvtFhuTMQM,9750
+vellum/workflows/resolvers/resolver.py,sha256=EIIiA2OlaVUPiKiSh6egQZxPA6ny1GDMdPq1AuN-mV8,2961
+vellum/workflows/resolvers/tests/test_resolver.py,sha256=PnUGzsulo1It_LjjhHsRNiILvvl5G_IaK8ZX56zKC28,6204
 vellum/workflows/resolvers/types.py,sha256=Hndhlk69g6EKLh_LYg5ILepW5U_h_BYNllfzhS9k8p4,237
 vellum/workflows/runner/__init__.py,sha256=i1iG5sAhtpdsrlvwgH6B-m49JsINkiWyPWs8vyT-bqM,72
 vellum/workflows/runner/runner.py,sha256=rwMZlQBkzK-EfewC6OysNCeuDeTCfpNuzaUlgYoFJMw,43329
 vellum/workflows/sandbox.py,sha256=mezSZmilR_fwR8164n8CEfzlMeQ55IqfapHp4ftImvQ,3212
 vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
-vellum/workflows/state/base.py,sha256=m9fCqbZn21GshCVCjJTD1dPZEQjFrsMXqlg7tM9fIwM,24283
+vellum/workflows/state/base.py,sha256=A8s0PC8UvFjPpkHDY6u-yIeb2KHjoAmu-GW-GYrDl0E,24654
 vellum/workflows/state/context.py,sha256=khM30U1iDNts5Xp8LXa_WfpkITNITexrDUUFJ5wZ2W4,8445
 vellum/workflows/state/delta.py,sha256=7h8wR10lRCm15SykaPj-gSEvvsMjCwYLPsOx3nsvBQg,440
 vellum/workflows/state/encoder.py,sha256=EynuS9aCt9Neb-H6HRCinEVZX5olCzME03W1TSXfpxs,1961
@@ -2037,8 +2037,8 @@ vellum/workflows/tests/test_sandbox.py,sha256=JKwaluI-lODQo7Ek9sjDstjL_WTdSqUlVi
 vellum/workflows/tests/test_undefined.py,sha256=zMCVliCXVNLrlC6hEGyOWDnQADJ2g83yc5FIM33zuo8,353
 vellum/workflows/types/__init__.py,sha256=KxUTMBGzuRCfiMqzzsykOeVvrrkaZmTTo1a7SLu8gRM,68
 vellum/workflows/types/code_execution_node_wrappers.py,sha256=fewX9bqF_4TZuK-gZYIn12s31-k03vHMGRpvFAPm11Y,3206
-vellum/workflows/types/core.py,sha256=yKm3sE02ult969q80DTmawiwYqodVjcAW-zlaUIgIv4,1495
-vellum/workflows/types/definition.py,sha256=PWJF1SAopdOZcLcxz1ZtGYOa0oH0-dxjFaa_tuGb6SI,7682
+vellum/workflows/types/core.py,sha256=X8xNBLUhnea97t3qeeZJ2XyUwW_a38YiznXr73j_ppE,950
+vellum/workflows/types/definition.py,sha256=Qof2MAjSNB0AN2XkSKmk-owuY59YcxDVHYpno6-StPA,8058
 vellum/workflows/types/generics.py,sha256=8jptbEx1fnJV0Lhj0MpCJOT6yNiEWeTOYOwrEAb5CRU,1576
 vellum/workflows/types/stack.py,sha256=h7NE0vXR7l9DevFBIzIAk1Zh59K-kECQtDTKOUunwMY,1314
 vellum/workflows/types/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -2065,8 +2065,8 @@ vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnad
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=Boa-_m9ii2Qsa1RvVM-VYniF7zCpzGgEGy-OnPZkrHg,23941
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-1.7.0.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-1.7.0.dist-info/METADATA,sha256=SMmQ-iMX4Dymys9Y_j7Jz4xq_QJGNUf8U9xYOgYKSw0,5547
-vellum_ai-1.7.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-1.7.0.dist-info/entry_points.txt,sha256=xVavzAKN4iF_NbmhWOlOkHluka0YLkbN_pFQ9pW3gLI,117
-vellum_ai-1.7.0.dist-info/RECORD,,
+vellum_ai-1.7.2.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-1.7.2.dist-info/METADATA,sha256=KyFKrTqcT-6Pvl5zTeo1EsEP0-uY4WREJPAr8Vdkh4A,5547
+vellum_ai-1.7.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-1.7.2.dist-info/entry_points.txt,sha256=xVavzAKN4iF_NbmhWOlOkHluka0YLkbN_pFQ9pW3gLI,117
+vellum_ai-1.7.2.dist-info/RECORD,,

vellum_ee/assets/node-definitions.json

@@ -699,6 +699,51 @@
         }
       ]
     },
+    {
+      "id": "ee9b5234-247b-49d1-bed4-490312f18838",
+      "display_data": {
+        "position": {
+          "x": 0.0,
+          "y": 0.0
+        },
+        "comment": {
+          "value": "Used to render a Jinja template.\n\n    Useful for lightweight data transformations and complex string templating.\n    ",
+          "expanded": true
+        }
+      },
+      "base": {
+        "name": "BaseNode",
+        "module": [
+          "vellum",
+          "workflows",
+          "nodes",
+          "bases",
+          "base"
+        ]
+      },
+      "definition": {
+        "name": "TemplatingNode",
+        "module": [
+          "vellum",
+          "workflows",
+          "nodes",
+          "core",
+          "templating_node",
+          "node"
+        ]
+      },
+      "trigger": {
+        "id": "005e1cd2-5452-4e1f-be6a-ac9fe3c02b9b",
+        "merge_behavior": "AWAIT_ATTRIBUTES"
+      },
+      "ports": [
+        {
+          "id": "8f4460f0-717b-4972-a6f4-ac164e5e204e",
+          "name": "default",
+          "type": "DEFAULT"
+        }
+      ]
+    },
     {
       "id": "035842e0-34ed-43af-97de-a706e40912ae",
       "label": "Tool Calling Node",
@@ -1010,10 +1055,5 @@
       ]
     }
   ],
-  "errors": [
-    {
-      "node": "TemplatingNode",
-      "error": "KeyError: TemplatingNode.Outputs.result"
-    }
-  ]
+  "errors": []
 }

vellum_ee/workflows/display/nodes/vellum/templating_node.py

@@ -50,7 +50,7 @@ class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_T
         # Misc type ignore is due to `node.Outputs` being generic
         # https://app.shortcut.com/vellum/story/4784
         output_descriptor = node.Outputs.result  # type: ignore [misc]
-        output_display = display_context.global_node_output_displays[output_descriptor]
+        output_display = self.get_node_output_display(output_descriptor)
         inferred_output_type = primitive_type_to_vellum_variable_type(output_descriptor)
 
         return {

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py

@@ -353,3 +353,33 @@ def test_serialize_node__adornment_order_matches_decorator_order():
     assert len(adornments) == 2
     assert adornments[0]["label"] == "Try Node"
     assert adornments[1]["label"] == "Retry Node"
+
+
+def test_serialize_workflow__retry_node_edges():
+    """
+    Tests that both retry-adorned nodes are correctly serialized in the nodes array.
+    """
+
+    @RetryNode.wrap(max_attempts=3, delay=60)
+    class FirstNode(BaseNode):
+        class Outputs(BaseOutputs):
+            value: str
+
+    @RetryNode.wrap(max_attempts=5, delay=120)
+    class SecondNode(BaseNode):
+        class Outputs(BaseOutputs):
+            result: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = FirstNode >> SecondNode
+
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    exec_config = cast(Dict[str, Any], workflow_display.serialize())
+
+    assert isinstance(exec_config["workflow_raw_data"], dict)
+    assert isinstance(exec_config["workflow_raw_data"]["nodes"], list)
+
+    nodes = cast(List[Dict[str, Any]], exec_config["workflow_raw_data"]["nodes"])
+
+    generic_nodes = [node for node in nodes if node["type"] == "GENERIC"]
+    assert len(generic_nodes) == 2

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -208,7 +208,15 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 self.display_context.add_invalid_node(node)
                 continue
 
-            serialized_nodes[node_display.node_id] = serialized_node
+            # Use wrapped node's ID as dict key for adornment wrappers to prevent overwrites
+            wrapped_node = get_wrapped_node(node)
+            if wrapped_node:
+                wrapped_node_display = self.display_context.node_displays[wrapped_node]
+                dict_key = wrapped_node_display.node_id
+            else:
+                dict_key = node_display.node_id
+
+            serialized_nodes[dict_key] = serialized_node
 
         synthetic_output_edges: JsonArray = []
         output_variables: JsonArray = []