vellum-ai 1.0.11__tar.gz → 1.1.1__tar.gz

vellum_ai-1.1.1/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.1
+Name: vellum-ai
+Version: 1.1.1
+Summary: 
+License: MIT
+Requires-Python: >=3.9,<4.0
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: OS Independent
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Dist: Jinja2 (>=3.1.0,<4.0.0)
+Requires-Dist: click (>=8.1.7,<9.0.0)
+Requires-Dist: docker (>=7.1.0,<8.0.0)
+Requires-Dist: httpx (>=0.21.2)
+Requires-Dist: openai (>=1.0.0,<2.0.0)
+Requires-Dist: orderly-set (>=5.2.2,<6.0.0)
+Requires-Dist: publication (==0.0.3)
+Requires-Dist: pydantic (>=1.9.2)
+Requires-Dist: pydantic-core (>=2.18.2,<3.0.0)
+Requires-Dist: pydash (>=7.0.0,<8.0.0)
+Requires-Dist: python-dateutil (>=2.8.0,<3.0.0)
+Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
+Requires-Dist: pytz (>=2022.0,<2026.0)
+Requires-Dist: pyyaml (>=6.0.0,<7.0.0)
+Requires-Dist: requests (>=2.31.0,<3.0.0)
+Requires-Dist: tomli (>=2.0.0,<3.0.0)
+Requires-Dist: typing_extensions (>=4.0.0)
+Project-URL: Repository, https://github.com/vellum-ai/vellum-python-sdks
+Description-Content-Type: text/markdown
+
+<p align="center">
+   <a href="https://vellum.ai">
+    <picture>
+      <img alt="Vellum README banner" src="https://storage.googleapis.com/vellum-public/assets/github-readme-banner.png" height="128">
+    </picture>
+  </a>
+  <p align="center">
+    <a href="https://vellum.ai">Learn more</a>
+    ·
+    <a href="https://www.vellum.ai/landing-pages/request-demo">Talk to us</a>
+  </p>
+</p>
+
+<p align="center">
+  <a href="https://pepy.tech/project/vellum-ai">
+    <img src="https://img.shields.io/pypi/dm/vellum-ai" alt="Vellum AI PyPI - Downloads" >
+  </a>
+  <a href="https://pypi.org/project/vellum-ai" target="_blank">
+    <img src="https://img.shields.io/pypi/v/vellum-ai?color=%2334D058&label=pypi%20package" alt="Package version">
+  </a>
+  <a href="https://www.ycombinator.com/companies/vellum">
+    <img src="https://img.shields.io/badge/Y%20Combinator-W23-orange?style=flat-square" alt="Y Combinator S24">
+  </a>
+</p>
+
+# Introduction
+
+[Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications.
+
+### Core Features
+
+- **Orchestration:** A powerful SDK and IDE for defining and debugging the control flow of your AI applications
+- **Prompting:** A best-in-class prompt playground for iterating on and refining prompts between models from any provider
+- **Evaluations**: An evaluations framework that makes it easy to measure the quality of your AI systems at scale
+- **Retrieval:** A ready-to-go service for turning unstructured content into intelligent, context-aware solutions
+  optimized for AI systems
+- **Deployment:** Decouple updates to your AI systems from your application code with an easy integration +
+  one-click deploy
+- **Observability:** Monitor and debug your AI systems in real-time with detailed logs, metrics, and end-user feedback
+
+## Table of Contents
+
+- [Get Started](#get-started)
+- [Client SDK](#client-sdk)
+- [Workflows SDK](#workflows-sdk)
+- [Contributing](#contributing)
+- [Open-source vs paid](#open-source-vs-paid)
+
+## Get Started
+
+Most functionality within the Vellum SDK requires a Vellum account and API key. You can sign up for free [here](https://app.vellum.ai/signup?f=wsdk&utm_source=github&utm_medium=repo_readme&utm_campaign=sdk) or visit our [pricing page](https://www.vellum.ai/pricing) for paid options.
+
+Even without a Vellum account, you can use the Workflows SDK to define the control flow of your AI systems. [Learn
+more below](#workflows-sdk).
+
+## Client SDK
+
+The Vellum Client SDK, found within `src/vellum/client` is a low-level client used to interact directly with the Vellum API.
+Learn more and get started by visiting the [Vellum Client SDK README](/src/vellum/client/README.md).
+
+## Workflows SDK
+
+The Vellum Workflows SDK is a high-level framework for defining and debugging the control flow of AI systems. At
+it's core, it's a powerful workflow engine with syntactic sugar for declaratively defining graphs, the nodes within,
+and the relationships between them.
+
+The Workflows SDK can be used with or without a Vellum account, but a Vellum account is required to use certain
+out-of-box nodes and features, including the ability to push and pull your Workflow definition to Vellum for editing
+and debugging via a UI.
+
+To learn more and get started, visit the [Vellum Workflows SDK README](/src/vellum/workflows/README.md).
+
+## Contributing
+
+See the [CONTRIBUTING.md](/CONTRIBUTING.md) for information on how to contribute to the Vellum SDKs.
+
+## Open-Source vs. Paid
+
+This repo is available under the [MIT expat license](https://github.com/vellum-ai/vellum-python-sdks/blob/main/LICENSE), except
+for the `ee` directory (which has its [license here](https://github.com/vellum-ai/vellum-python-sdks/blob/main/ee/LICENSE)) if applicable.
+
+To learn more, [book a demo](https://www.vellum.ai/landing-pages/request-demo) or see our [pricing page](https://www.vellum.ai/pricing).
+

vellum_ai-1.1.1/ee/vellum_ee/workflows/display/exceptions.py

@@ -0,0 +1,7 @@
+class NodeValidationError(Exception):
+    """Exception raised when a node fails validation during workflow display serialization."""
+
+    def __init__(self, message: str, node_class_name: str):
+        self.message = message
+        self.node_class_name = node_class_name
+        super().__init__(f"Node validation error in {node_class_name}: {message}")

vellum_ai-1.1.1/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -0,0 +1,110 @@
+import inspect
+from uuid import UUID
+from typing import ClassVar, Generic, Optional, TypeVar
+
+from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
+from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
+from vellum.workflows.types.core import JsonObject
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+from vellum_ee.workflows.display.exceptions import NodeValidationError
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+_CodeExecutionNodeType = TypeVar("_CodeExecutionNodeType", bound=CodeExecutionNode)
+
+
+class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
+    output_id: ClassVar[Optional[UUID]] = None
+    log_output_id: ClassVar[Optional[UUID]] = None
+
+    __serializable_inputs__ = {
+        CodeExecutionNode.code,
+        CodeExecutionNode.code_inputs,
+    }
+
+    def serialize(
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
+    ) -> JsonObject:
+        node = self._node
+        node_id = self.node_id
+        raw_code = raise_if_descriptor(node.code)
+        filepath = raise_if_descriptor(node.filepath)
+
+        code_value: Optional[str]
+        if raw_code:
+            code_value = raw_code
+        elif filepath:
+            node_file_path = inspect.getfile(node)
+            file_code = read_file_from_path(
+                node_filepath=node_file_path,
+                script_filepath=filepath,
+            )
+            if not file_code:
+                raise NodeValidationError(f"Filepath '{filepath}' does not exist", node.__name__)
+            code_value = file_code
+        else:
+            code_value = ""
+
+        code_inputs = raise_if_descriptor(node.code_inputs)
+
+        inputs = [
+            create_node_input(
+                node_id=node_id,
+                input_name=variable_name,
+                value=variable_value,
+                display_context=display_context,
+                input_id=self.node_input_ids_by_name.get(f"{CodeExecutionNode.code_inputs.name}.{variable_name}")
+                or self.node_input_ids_by_name.get(variable_name),
+            )
+            for variable_name, variable_value in code_inputs.items()
+        ]
+
+        code_input_id = self.node_input_ids_by_name.get(CodeExecutionNode.code.name)
+        code_node_input = create_node_input(
+            node_id=node_id,
+            input_name="code",
+            value=code_value,
+            display_context=display_context,
+            input_id=code_input_id,
+        )
+
+        runtime_input_id = self.node_input_ids_by_name.get(CodeExecutionNode.runtime.name)
+        runtime_node_input = create_node_input(
+            node_id=node_id,
+            input_name="runtime",
+            value=node.runtime,
+            display_context=display_context,
+            input_id=runtime_input_id,
+        )
+        inputs.extend([code_node_input, runtime_node_input])
+
+        packages = raise_if_descriptor(node.packages)
+
+        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())
+
+        return {
+            "id": str(node_id),
+            "type": "CODE_EXECUTION",
+            "inputs": [input.dict() for input in inputs],
+            "data": {
+                "label": self.label,
+                "error_output_id": str(error_output_id) if error_output_id else None,
+                "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
+                "target_handle_id": str(self.get_target_handle_id()),
+                "code_input_id": code_node_input.id,
+                "runtime_input_id": runtime_node_input.id,
+                "output_type": output_type,
+                "packages": [package.dict() for package in packages] if packages is not None else [],
+                "output_id": str(self.output_id) if self.output_id else str(output_display.id),
+                "log_output_id": str(self.log_output_id) if self.log_output_id else str(log_output_display.id),
+            },
+            "display_data": self.get_display_data().dict(),
+            "base": self.get_base().dict(),
+            "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
+        }

vellum_ai-1.1.1/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py

@@ -0,0 +1,222 @@
+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.exceptions import NodeValidationError
+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
+
+
+def _no_display_class(Node: Type[CodeExecutionNode]):
+    return None
+
+
+def _display_class_with_node_input_ids_by_name(Node: Type[CodeExecutionNode]):
+    class CodeExecutionNodeDisplay(BaseCodeExecutionNodeDisplay[Node]):  # type: ignore[valid-type]
+        node_input_ids_by_name = {"foo": UUID("fba6a4d5-835a-4e99-afb7-f6a4aed15110")}
+
+    return CodeExecutionNodeDisplay
+
+
+def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[CodeExecutionNode]):
+    class CodeExecutionNodeDisplay(BaseCodeExecutionNodeDisplay[Node]):  # type: ignore[valid-type]
+        node_input_ids_by_name = {"code_inputs.foo": UUID("fba6a4d5-835a-4e99-afb7-f6a4aed15110")}
+
+    return CodeExecutionNodeDisplay
+
+
+@pytest.mark.parametrize(
+    ["GetDisplayClass", "expected_input_id"],
+    [
+        (_no_display_class, "a5dbe403-0b00-4df6-b8f7-ed5f7794b003"),
+        (_display_class_with_node_input_ids_by_name, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
+        (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
+    ],
+    ids=[
+        "no_display_class",
+        "display_class_with_node_input_ids_by_name",
+        "display_class_with_node_input_ids_by_name_with_inputs_prefix",
+    ],
+)
+def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
+    # GIVEN a code node with inputs
+    class MyCodeExecutionNode(CodeExecutionNode):
+        code_inputs = {"foo": "bar"}
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyCodeExecutionNode
+
+    # AND a display class
+    GetDisplayClass(MyCodeExecutionNode)
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the inputs
+    my_code_execution_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "CODE_EXECUTION"
+    )
+
+    assert my_code_execution_node["inputs"] == [
+        {
+            "id": expected_input_id,
+            "key": "foo",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "bar",
+                        },
+                    }
+                ],
+            },
+        },
+        {
+            "id": "9774d864-c76d-4a1a-8181-b632ed3ab87c",
+            "key": "code",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "",
+                        },
+                    }
+                ],
+            },
+        },
+        {
+            "id": "34742235-5699-45cd-9d34-bce3745e743d",
+            "key": "runtime",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "PYTHON_3_11_6",
+                        },
+                    }
+                ],
+            },
+        },
+    ]
+
+
+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])
+
+
+def test_serialize_node__with_non_exist_code_input_path():
+    # GIVEN a code node with a non-existent code input path
+    class MyNode(CodeExecutionNode):
+        filepath = "non_existent_file.py"
+
+    # 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)
+    with pytest.raises(NodeValidationError) as exc_info:
+        workflow_display.serialize()
+    assert "Filepath 'non_existent_file.py' does not exist" in str(exc_info.value)
+
+
+def test_serialize_node__with_non_exist_code_input_path_with_dry_run():
+    # GIVEN a code node with a non-existent code input path
+    class MyNode(CodeExecutionNode):
+        filepath = "non_existent_file.py"
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow with dry_run=True
+    workflow_display = get_workflow_display(workflow_class=Workflow, dry_run=True)
+    data: dict = workflow_display.serialize()
+
+    # THEN the workflow should not raise an error
+    assert data == {
+        "workflow_raw_data": {
+            "nodes": [
+                {
+                    "id": "9b9e2a5d-01a4-46b2-80a3-d9484b2c0e08",
+                    "type": "ENTRYPOINT",
+                    "inputs": [],
+                    "data": {"label": "Entrypoint Node", "source_handle_id": "3e2a3f52-5047-4e2e-9a21-37bd43c63250"},
+                    "display_data": {"position": {"x": 0.0, "y": -50.0}},
+                    "base": None,
+                    "definition": None,
+                }
+            ],
+            "edges": [
+                {
+                    "id": "ab6ef06e-df2c-4877-9c3e-9d7261b39748",
+                    "source_node_id": "9b9e2a5d-01a4-46b2-80a3-d9484b2c0e08",
+                    "source_handle_id": "3e2a3f52-5047-4e2e-9a21-37bd43c63250",
+                    "target_node_id": "ac90c0ce-f393-438c-a24f-e5e9a9286182",
+                    "target_handle_id": "3a39ea63-9f86-4891-a902-0216a7190720",
+                    "type": "DEFAULT",
+                }
+            ],
+            "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
+            "definition": {
+                "name": "Workflow",
+                "module": ["vellum_ee", "workflows", "display", "nodes", "vellum", "tests", "test_code_execution_node"],
+            },
+            "output_values": [],
+        },
+        "input_variables": [],
+        "state_variables": [],
+        "output_variables": [],
+    }

vellum_ai-1.1.1/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -0,0 +1,217 @@
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.definition import AuthorizationType, EnvironmentVariableReference, MCPServer
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_serialize_node__prompt_inputs__constant_value():
+    # GIVEN a prompt node with constant value inputs
+    class MyPromptNode(ToolCallingNode):
+        prompt_inputs = {"foo": "bar"}
+
+    # AND a workflow with the prompt node
+    class Workflow(BaseWorkflow):
+        graph = MyPromptNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the inputs
+    my_prompt_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
+    )
+
+    prompt_inputs_attribute = next(
+        attribute for attribute in my_prompt_node["attributes"] if attribute["name"] == "prompt_inputs"
+    )
+
+    assert prompt_inputs_attribute == {
+        "id": "3d9a4d2e-c9bd-4417-8a0c-52f15efdbe30",
+        "name": "prompt_inputs",
+        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": {"foo": "bar"}}},
+    }
+
+
+def test_serialize_node__prompt_inputs__input_reference():
+    # GIVEN a state definition
+    class MyInput(BaseInputs):
+        foo: str
+
+    # AND a prompt node with inputs
+    class MyPromptNode(InlinePromptNode):
+        prompt_inputs = {"foo": MyInput.foo}
+
+    # AND a workflow with the prompt node
+    class Workflow(BaseWorkflow[MyInput, BaseState]):
+        graph = MyPromptNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should skip the state reference input rule
+    my_prompt_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
+    )
+
+    prompt_inputs_attribute = next(
+        attribute for attribute in my_prompt_node["attributes"] if attribute["name"] == "prompt_inputs"
+    )
+
+    assert prompt_inputs_attribute == {
+        "id": "6cde4776-7f4a-411c-95a8-69c8b3a64b42",
+        "name": "prompt_inputs",
+        "value": {
+            "type": "DICTIONARY_REFERENCE",
+            "entries": [
+                {
+                    "id": "ab7902ef-de14-4edc-835c-366d3ef6a70e",
+                    "key": "foo",
+                    "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "e3657390-fd3c-4fea-8cdd-fc5ea79f3278"},
+                }
+            ],
+        },
+    }
+
+
+def test_serialize_node__prompt_inputs__mixed_values():
+    # GIVEN a prompt node with mixed values
+    class MyInput(BaseInputs):
+        foo: str
+
+    # AND a prompt node with mixed values
+    class MyPromptNode(InlinePromptNode):
+        prompt_inputs = {"foo": "bar", "baz": MyInput.foo}
+
+    # AND a workflow with the prompt node
+    class Workflow(BaseWorkflow[MyInput, BaseState]):
+        graph = MyPromptNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the inputs
+    my_prompt_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
+    )
+
+    prompt_inputs_attribute = next(
+        attribute for attribute in my_prompt_node["attributes"] if attribute["name"] == "prompt_inputs"
+    )
+
+    assert prompt_inputs_attribute == {
+        "id": "c4ca6e3d-0f71-4802-a618-1e87880cb7cf",
+        "name": "prompt_inputs",
+        "value": {
+            "type": "DICTIONARY_REFERENCE",
+            "entries": [
+                {
+                    "id": "0fc7e25e-075c-4849-b89b-9729d1aeada1",
+                    "key": "foo",
+                    "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "bar"}},
+                },
+                {
+                    "id": "bba42c89-fa7b-4cb7-bc16-0d21ce060a4b",
+                    "key": "baz",
+                    "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "8d57cf1d-147c-427b-9a5e-e5f6ab76e2eb"},
+                },
+            ],
+        },
+    }
+
+
+def test_serialize_node__tool_calling_node__mcp_server_api_key():
+    # GIVEN a tool calling node with an mcp server
+    class MyToolCallingNode(ToolCallingNode):
+        functions = [
+            MCPServer(
+                name="my-mcp-server",
+                url="https://my-mcp-server.com",
+                authorization_type=AuthorizationType.API_KEY,
+                api_key_header_key="my-api-key-header-key",
+                api_key_header_value=EnvironmentVariableReference(name="my-api-key-header-value"),
+            )
+        ]
+
+    # AND a workflow with the tool calling node
+    class Workflow(BaseWorkflow):
+        graph = MyToolCallingNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the mcp server
+    my_tool_calling_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(MyToolCallingNode.__id__)
+    )
+
+    functions_attribute = next(
+        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
+    )
+
+    assert functions_attribute == {
+        "id": "6c0f7d4f-3c8a-4201-b588-8398d3c97480",
+        "name": "functions",
+        "value": {
+            "type": "ARRAY_REFERENCE",
+            "items": [
+                {
+                    "type": "DICTIONARY_REFERENCE",
+                    "entries": [
+                        {
+                            "id": "bcf2713b-19fc-4b4b-8ff5-b45c8e63c665",
+                            "key": "type",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "MCP_SERVER"}},
+                        },
+                        {
+                            "id": "4e00439e-ce6f-4e0a-be4c-0fc05990ec44",
+                            "key": "name",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "my-mcp-server"}},
+                        },
+                        {
+                            "id": "9a3dcca9-4595-4efb-ada6-c011721f7018",
+                            "key": "url",
+                            "value": {
+                                "type": "CONSTANT_VALUE",
+                                "value": {"type": "STRING", "value": "https://my-mcp-server.com"},
+                            },
+                        },
+                        {
+                            "id": "21c3a6ce-7607-42a4-92b8-9d7c4061edce",
+                            "key": "authorization_type",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "API_KEY"}},
+                        },
+                        {
+                            "id": "dcf8e8f0-84d2-4ffb-b0e6-218c569015a4",
+                            "key": "bearer_token_value",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "fe93f516-42c0-40cf-b476-50bd04857d5f",
+                            "key": "api_key_header_key",
+                            "value": {
+                                "type": "CONSTANT_VALUE",
+                                "value": {"type": "STRING", "value": "my-api-key-header-key"},
+                            },
+                        },
+                        {
+                            "id": "750ed8ca-3bf3-46be-96f0-2c92d19d2084",
+                            "key": "api_key_header_value",
+                            "value": {
+                                "type": "ENVIRONMENT_VARIABLE",
+                                "environment_variable": "my-api-key-header-value",
+                            },
+                        },
+                    ],
+                }
+            ],
+        },
+    }