vellum-ai 0.14.56__tar.gz → 0.14.57__tar.gz

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/PKG-INFO

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

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

@@ -0,0 +1,118 @@
+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.experimental.tool_calling_node.node import ToolCallingNode
+from vellum.workflows.state.base import BaseState
+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": [
+                {
+                    "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": [
+                {"key": "foo", "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "bar"}}},
+                {
+                    "key": "baz",
+                    "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "8d57cf1d-147c-427b-9a5e-e5f6ab76e2eb"},
+                },
+            ],
+        },
+    }

vellum_ai-0.14.57/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py

@@ -0,0 +1,275 @@
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_inline_prompt_node_with_functions.workflow import BasicInlinePromptWithFunctionsWorkflow
+
+
+def test_serialize_workflow():
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=BasicInlinePromptWithFunctionsWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "ceb5cc94-48ee-4968-b37a-421623a8f1ef",
+                "key": "noun",
+                "type": "STRING",
+                "default": None,
+                "required": True,
+                "extensions": {"color": None},
+            }
+        ],
+        input_variables,
+        ignore_order=True,
+    )
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert not DeepDiff(
+        [{"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "JSON"}],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert len(workflow_raw_data["edges"]) == 2
+    assert len(workflow_raw_data["nodes"]) == 3
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "382842a3-0490-4dee-b87b-eef86766f07c",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {"label": "Entrypoint Node", "source_handle_id": "8294baa6-8bf4-4b54-a56b-407b64851b77"},
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "base": None,
+        "definition": None,
+    }
+
+    prompt_node = workflow_raw_data["nodes"][1]
+    assert not DeepDiff(
+        {
+            "id": "8450dd06-975a-41a4-a564-808ee8808fe6",
+            "type": "PROMPT",
+            "inputs": [
+                {
+                    "id": "f7fca55e-93e9-4009-9227-acf839c7990d",
+                    "key": "noun",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "INPUT_VARIABLE",
+                                "data": {"input_variable_id": "ceb5cc94-48ee-4968-b37a-421623a8f1ef"},
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "data": {
+                "label": "Example Base Inline Prompt Node With Functions",
+                "output_id": "ead0ccb5-092f-4d9b-a9ec-5eb83d498188",
+                "error_output_id": None,
+                "array_output_id": "628df199-a049-40b9-a29b-a378edd759bb",
+                "source_handle_id": "d4a097ab-e22d-42f1-b6bc-2ed96856377a",
+                "target_handle_id": "c2dccecb-8a41-40a8-95af-325d3ab8bfe5",
+                "variant": "INLINE",
+                "exec_config": {
+                    "parameters": {
+                        "stop": [],
+                        "temperature": 0.0,
+                        "max_tokens": 4096,
+                        "top_p": 1.0,
+                        "top_k": 0,
+                        "frequency_penalty": 0.0,
+                        "presence_penalty": 0.0,
+                        "logit_bias": None,
+                        "custom_parameters": None,
+                    },
+                    "input_variables": [
+                        {"id": "f7fca55e-93e9-4009-9227-acf839c7990d", "key": "noun", "type": "STRING"}
+                    ],
+                    "prompt_template_block_data": {
+                        "version": 1,
+                        "blocks": [
+                            {
+                                "block_type": "CHAT_MESSAGE",
+                                "properties": {
+                                    "chat_role": "SYSTEM",
+                                    "chat_source": None,
+                                    "chat_message_unterminated": False,
+                                    "blocks": [
+                                        {
+                                            "block_type": "JINJA",
+                                            "properties": {
+                                                "template": "What's your favorite {{noun}}?",
+                                                "template_type": "STRING",
+                                            },
+                                            "id": "467fe2b1-312b-40db-8869-9c6ada7c7077",
+                                            "cache_config": None,
+                                            "state": "ENABLED",
+                                        }
+                                    ],
+                                },
+                                "id": "1d1e117d-19dc-4282-b1e3-9534014fb6e5",
+                                "cache_config": None,
+                                "state": "ENABLED",
+                            },
+                            {
+                                "id": "9b34f084-449d-423f-8691-37518b1ee9ca",
+                                "block_type": "FUNCTION_DEFINITION",
+                                "properties": {
+                                    "function_name": "favorite_noun",
+                                    "function_description": "Returns the favorite noun of the user",
+                                    "function_parameters": {},
+                                    "function_forced": None,
+                                    "function_strict": None,
+                                },
+                            },
+                        ],
+                    },
+                },
+                "ml_model_name": "gpt-4o",
+            },
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {
+                "name": "InlinePromptNode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "inline_prompt_node", "node"],
+            },
+            "definition": {
+                "name": "ExampleBaseInlinePromptNodeWithFunctions",
+                "module": ["tests", "workflows", "basic_inline_prompt_node_with_functions", "workflow"],
+            },
+            "outputs": [
+                {"id": "9557bd86-702d-4b45-b8c1-c3980bffe28f", "name": "json", "type": "JSON", "value": None},
+                {"id": "ead0ccb5-092f-4d9b-a9ec-5eb83d498188", "name": "text", "type": "STRING", "value": None},
+                {"id": "628df199-a049-40b9-a29b-a378edd759bb", "name": "results", "type": "ARRAY", "value": None},
+            ],
+            "ports": [{"id": "d4a097ab-e22d-42f1-b6bc-2ed96856377a", "name": "default", "type": "DEFAULT"}],
+            "attributes": [
+                {
+                    "id": "6cd5395c-6e46-4bc9-b98c-8f8924554555",
+                    "name": "ml_model",
+                    "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "gpt-4o"}},
+                },
+                {
+                    "id": "ffabe7d2-8ab6-4201-9d41-c4d7be1386e1",
+                    "name": "prompt_inputs",
+                    "value": {
+                        "type": "DICTIONARY_REFERENCE",
+                        "entries": [
+                            {
+                                "key": "noun",
+                                "value": {
+                                    "type": "WORKFLOW_INPUT",
+                                    "input_variable_id": "ceb5cc94-48ee-4968-b37a-421623a8f1ef",
+                                },
+                            }
+                        ],
+                    },
+                },
+            ],
+        },
+        prompt_node,
+        ignore_order=True,
+    )
+
+    final_output_node = workflow_raw_data["nodes"][2]
+    assert not DeepDiff(
+        {
+            "id": "42318326-3ae8-417f-9609-f6d8ae47eafb",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "results",
+                "target_handle_id": "46c99277-2b4b-477d-851c-ea497aef6b16",
+                "output_id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e",
+                "output_type": "JSON",
+                "node_input_id": "d7c89dce-765b-494d-a256-aba4bcf87b42",
+            },
+            "inputs": [
+                {
+                    "id": "d7c89dce-765b-494d-a256-aba4bcf87b42",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "8450dd06-975a-41a4-a564-808ee8808fe6",
+                                    "output_id": "628df199-a049-40b9-a29b-a378edd759bb",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {
+                "name": "FinalOutputNode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
+            },
+            "definition": None,
+        },
+        final_output_node,
+        ignore_order=True,
+    )
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "924f693f-3f4c-466a-8cde-648ba3baf9fd",
+                "source_node_id": "382842a3-0490-4dee-b87b-eef86766f07c",
+                "source_handle_id": "8294baa6-8bf4-4b54-a56b-407b64851b77",
+                "target_node_id": "8450dd06-975a-41a4-a564-808ee8808fe6",
+                "target_handle_id": "c2dccecb-8a41-40a8-95af-325d3ab8bfe5",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "05ca58fb-e02d-48d4-9207-2dad0833a25b",
+                "source_node_id": "8450dd06-975a-41a4-a564-808ee8808fe6",
+                "source_handle_id": "d4a097ab-e22d-42f1-b6bc-2ed96856377a",
+                "target_node_id": "42318326-3ae8-417f-9609-f6d8ae47eafb",
+                "target_handle_id": "46c99277-2b4b-477d-851c-ea497aef6b16",
+                "type": "DEFAULT",
+            },
+        ],
+        serialized_edges,
+        ignore_order=True,
+    )
+
+    # AND the display data should be what we expect
+    display_data = workflow_raw_data["display_data"]
+    assert display_data == {
+        "viewport": {
+            "x": 0.0,
+            "y": 0.0,
+            "zoom": 1.0,
+        }
+    }
+
+    # AND the definition should be what we expect
+    definition = workflow_raw_data["definition"]
+    assert definition == {
+        "name": "BasicInlinePromptWithFunctionsWorkflow",
+        "module": ["tests", "workflows", "basic_inline_prompt_node_with_functions", "workflow"],
+    }

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py

@@ -130,11 +130,12 @@ def test_serialize_workflow():
                         "type": "JSON",
                         "value": [
                             {
+                                "type": "CODE_EXECUTION",
                                 "definition": {
                                     "state": None,
                                     "cache_config": None,
                                     "name": "get_current_weather",
-                                    "description": None,
+                                    "description": "\n    Get the current weather in a given location.\n    ",
                                     "parameters": {
                                         "type": "object",
                                         "properties": {"location": {"type": "string"}, "unit": {"type": "string"}},

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "0.14.56"
+version = "0.14.57"
 description = ""
 readme = "README.md"
 authors = []

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

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/src/vellum/workflows/nodes/bases/base.py

@@ -8,7 +8,6 @@ from typing import Any, Dict, Generic, Iterator, Optional, Set, Tuple, Type, Typ
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.descriptors.utils import is_unresolved, resolve_value
-from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.graph import Graph
@@ -317,8 +316,15 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
                 """
                 # Check if all dependencies have invoked this node
                 dependencies_invoked = state.meta.node_execution_cache._dependencies_invoked.get(node_span_id, set())
-                all_deps_invoked = all(dep in dependencies_invoked for dep in dependencies)
-
+                node_classes_invoked = {
+                    state.meta.node_execution_cache.__node_execution_lookup__[dep]
+                    for dep in dependencies_invoked
+                    if dep in state.meta.node_execution_cache.__node_execution_lookup__
+                }
+                if len(node_classes_invoked) != len(dependencies):
+                    return False
+
+                all_deps_invoked = all(dep in node_classes_invoked for dep in dependencies)
                 return all_deps_invoked
 
             raise NodeException(
@@ -328,7 +334,7 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
 
         @classmethod
         def _queue_node_execution(
-            cls, state: StateType, dependencies: Set["Type[BaseNode]"], invoked_by: Optional[Edge] = None
+            cls, state: StateType, dependencies: Set["Type[BaseNode]"], invoked_by: Optional[UUID] = None
         ) -> UUID:
             """
             Queues a future execution of a node, returning the span id of the execution.
@@ -341,19 +347,21 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
             if not invoked_by:
                 return execution_id
 
+            if invoked_by not in state.meta.node_execution_cache.__node_execution_lookup__:
+                return execution_id
+
             if cls.merge_behavior not in {MergeBehavior.AWAIT_ANY, MergeBehavior.AWAIT_ALL}:
                 return execution_id
 
-            source_node = invoked_by.from_port.node_class
             for queued_node_execution_id in state.meta.node_execution_cache._node_executions_queued[cls.node_class]:
-                if source_node not in state.meta.node_execution_cache._dependencies_invoked[queued_node_execution_id]:
+                if invoked_by not in state.meta.node_execution_cache._dependencies_invoked[queued_node_execution_id]:
                     state.meta.node_execution_cache._invoke_dependency(
-                        queued_node_execution_id, cls.node_class, source_node, dependencies
+                        queued_node_execution_id, cls.node_class, invoked_by, dependencies
                     )
                     return queued_node_execution_id
 
             state.meta.node_execution_cache._node_executions_queued[cls.node_class].append(execution_id)
-            state.meta.node_execution_cache._invoke_dependency(execution_id, cls.node_class, source_node, dependencies)
+            state.meta.node_execution_cache._invoke_dependency(execution_id, cls.node_class, invoked_by, dependencies)
             return execution_id
 
     class Execution(metaclass=_BaseNodeExecutionMeta):

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/src/vellum/workflows/nodes/core/retry_node/node.py

@@ -70,6 +70,12 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
 
                     for output_descriptor, output_value in event.outputs:
                         setattr(node_outputs, output_descriptor.name, output_value)
+
+                        if self.__wrapped_node__:
+                            inner_desc = getattr(self.__wrapped_node__.Outputs, output_descriptor.name, None)
+                            if inner_desc:
+                                self.state.meta.node_outputs[inner_desc] = output_value
+
                 elif event.name == "workflow.execution.paused":
                     exception = NodeException(
                         code=WorkflowErrorCode.INVALID_OUTPUTS,

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/src/vellum/workflows/nodes/displayable/api_node/node.py

@@ -47,11 +47,18 @@ class APINode(BaseAPINode):
             self.bearer_token_value, VellumSecret
         ):
             bearer_token = self.bearer_token_value
+
+        final_headers = {**headers, **header_overrides}
+
+        vellum_client_wrapper = self._context.vellum_client._client_wrapper
+        if self.url.startswith(vellum_client_wrapper._environment.default) and "X-API-Key" not in final_headers:
+            final_headers["X-API-Key"] = vellum_client_wrapper.api_key
+
         return self._run(
             method=self.method,
             url=self.url,
             data=self.data,
             json=self.json,
-            headers={**headers, **header_overrides},
+            headers=final_headers,
             bearer_token=bearer_token,
         )

{vellum_ai-0.14.56 → vellum_ai-0.14.57}/src/vellum/workflows/nodes/displayable/api_node/tests/test_api_node.py

@@ -19,7 +19,7 @@ def test_run_workflow__secrets(vellum_client):
     class SimpleBaseAPINode(APINode):
         method = APIRequestMethod.POST
         authorization_type = AuthorizationType.BEARER_TOKEN
-        url = "https://api.vellum.ai"
+        url = "https://example.vellum.ai"
         json = {
             "key": "value",
         }
@@ -32,7 +32,9 @@ def test_run_workflow__secrets(vellum_client):
     terminal = node.run()
 
     assert vellum_client.execute_api.call_count == 1
+    assert vellum_client.execute_api.call_args.kwargs["url"] == "https://example.vellum.ai"
     assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
+    assert vellum_client.execute_api.call_args.kwargs["headers"] == {"X-Test-Header": "foo"}
     bearer_token = vellum_client.execute_api.call_args.kwargs["bearer_token"]
     assert bearer_token == ClientVellumSecret(name="secret")
     assert terminal.headers == {"X-Response-Header": "bar"}
@@ -45,7 +47,7 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
     class SimpleAPINode(APINode):
         method = APIRequestMethod.GET
         authorization_type = AuthorizationType.BEARER_TOKEN
-        url = "https://api.vellum.ai"
+        url = "https://example.vellum.ai"
         json = {
             "key": "value",
         }
@@ -65,7 +67,9 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
 
     # AND the API call should have been made
     assert vellum_client.execute_api.call_count == 1
+    assert vellum_client.execute_api.call_args.kwargs["url"] == "https://example.vellum.ai"
     assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
+    assert vellum_client.execute_api.call_args.kwargs["headers"] == {"X-Test-Header": "foo"}
 
 
 def test_api_node_defaults_to_get_method(vellum_client):
@@ -80,7 +84,7 @@ def test_api_node_defaults_to_get_method(vellum_client):
     # AND an API node without a method specified
     class SimpleAPINodeWithoutMethod(APINode):
         authorization_type = AuthorizationType.BEARER_TOKEN
-        url = "https://api.vellum.ai"
+        url = "https://example.vellum.ai"
         headers = {
             "X-Test-Header": "foo",
         }
@@ -95,3 +99,62 @@ def test_api_node_defaults_to_get_method(vellum_client):
     assert vellum_client.execute_api.call_count == 1
     method = vellum_client.execute_api.call_args.kwargs["method"]
     assert method == APIRequestMethod.GET.value
+
+
+def test_api_node__detects_client_environment_urls__adds_token(mock_httpx_transport, mock_requests, monkeypatch):
+    # GIVEN an API node with a URL pointing back to Vellum
+    class SimpleAPINodeToVellum(APINode):
+        url = "https://api.vellum.ai"
+
+    # AND a mock request sent to the Vellum API would return a 200
+    mock_response = mock_requests.get(
+        "https://api.vellum.ai",
+        status_code=200,
+        json={"data": [1, 2, 3]},
+    )
+
+    # AND an api key is set
+    monkeypatch.setenv("VELLUM_API_KEY", "vellum-api-key-1234")
+
+    # WHEN we run the node
+    node = SimpleAPINodeToVellum()
+    node.run()
+
+    # THEN the execute_api method should not have been called
+    mock_httpx_transport.handle_request.assert_not_called()
+
+    # AND the vellum API should have been called with the correct headers
+    assert mock_response.last_request
+    assert mock_response.last_request.headers["X-API-Key"] == "vellum-api-key-1234"
+
+
+def test_api_node__detects_client_environment_urls__does_not_override_headers(
+    mock_httpx_transport, mock_requests, monkeypatch
+):
+    # GIVEN an API node with a URL pointing back to Vellum
+    class SimpleAPINodeToVellum(APINode):
+        url = "https://api.vellum.ai"
+        headers = {
+            "X-API-Key": "vellum-api-key-5678",
+        }
+
+    # AND a mock request sent to the Vellum API would return a 200
+    mock_response = mock_requests.get(
+        "https://api.vellum.ai",
+        status_code=200,
+        json={"data": [1, 2, 3]},
+    )
+
+    # AND an api key is set
+    monkeypatch.setenv("VELLUM_API_KEY", "vellum-api-key-1234")
+
+    # WHEN we run the node
+    node = SimpleAPINodeToVellum()
+    node.run()
+
+    # THEN the execute_api method should not have been called
+    mock_httpx_transport.handle_request.assert_not_called()
+
+    # AND the vellum API should have been called with the correct headers
+    assert mock_response.last_request
+    assert mock_response.last_request.headers["X-API-Key"] == "vellum-api-key-5678"