vellum-ai 1.0.6__tar.gz → 1.0.8__tar.gz

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/PKG-INFO

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

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/ee/vellum_cli/push.py

@@ -95,8 +95,9 @@ def push_command(
         and workflow_config.workflow_sandbox_id
         and workflow_config.workflow_sandbox_id != workflow_sandbox_id
     ):
-        raise ValueError(
-            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'."
+        logger.warning(
+            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'. "
+            f"Continuing with the provided workflow sandbox id '{workflow_sandbox_id}'."
         )
 
     client = create_vellum_client(

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py

@@ -140,13 +140,19 @@ def test_serialize_workflow():
                                     "description": "\n    Get the current weather in a given location.\n    ",
                                     "parameters": {
                                         "type": "object",
-                                        "properties": {"location": {"type": "string"}, "unit": {"type": "string"}},
+                                        "properties": {
+                                            "location": {
+                                                "type": "string",
+                                                "description": "The location to get the weather for",
+                                            },
+                                            "unit": {"type": "string", "description": "The unit of temperature"},
+                                        },
                                         "required": ["location", "unit"],
                                     },
                                     "forced": None,
                                     "strict": None,
                                 },
-                                "src": 'import math\n\n\ndef get_current_weather(location: str, unit: str) -> str:\n    """\n    Get the current weather in a given location.\n    """\n    return f"The current weather in {location} is sunny with a temperature of {get_temperature(70.1)} degrees {unit}."\n\n\ndef get_temperature(temperature: float) -> int:\n    """\n    Get the temperature in a given location.\n    """\n    return math.floor(temperature)\n',  # noqa: E501
+                                "src": 'import math\nfrom typing import Annotated\n\n\ndef get_current_weather(\n    location: Annotated[str, "The location to get the weather for"], unit: Annotated[str, "The unit of temperature"]\n) -> str:\n    """\n    Get the current weather in a given location.\n    """\n    return f"The current weather in {location} is sunny with a temperature of {get_temperature(70.1)} degrees {unit}."\n\n\ndef get_temperature(temperature: float) -> int:\n    """\n    Get the temperature in a given location.\n    """\n    return math.floor(temperature)\n',  # noqa: E501
                             }
                         ],
                     },

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/pyproject.toml

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

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/src/vellum/client/core/client_wrapper.py

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

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

@@ -16,6 +16,7 @@ def test_run_workflow__secrets(vellum_client):
         json_={"data": [1, 2, 3]},
         headers={"X-Response-Header": "bar"},
     )
+    vellum_client._client_wrapper.get_headers.return_value = {"User-Agent": "vellum-ai/1.0.6"}
 
     class SimpleBaseAPINode(APINode):
         method = APIRequestMethod.POST
@@ -35,7 +36,9 @@ def test_run_workflow__secrets(vellum_client):
     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"}
+    headers = vellum_client.execute_api.call_args.kwargs["headers"]
+    assert headers["X-Test-Header"] == "foo"
+    assert "vellum-ai" in headers.get("User-Agent", "")
     bearer_token = vellum_client.execute_api.call_args.kwargs["bearer_token"]
     assert bearer_token == ClientVellumSecret(name="secret")
     assert terminal.headers == {"X-Response-Header": "bar"}
@@ -44,6 +47,7 @@ def test_run_workflow__secrets(vellum_client):
 def test_api_node_raises_error_when_api_call_fails(vellum_client):
     # GIVEN an API call that fails
     vellum_client.execute_api.side_effect = ApiError(status_code=400, body="API Error")
+    vellum_client._client_wrapper.get_headers.return_value = {"User-Agent": "vellum-ai/1.0.6"}
 
     class SimpleAPINode(APINode):
         method = APIRequestMethod.GET
@@ -70,7 +74,9 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
     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"}
+    headers = vellum_client.execute_api.call_args.kwargs["headers"]
+    assert headers["X-Test-Header"] == "foo"
+    assert "vellum-ai" in headers.get("User-Agent", "")
 
 
 def test_api_node_defaults_to_get_method(vellum_client):

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/src/vellum/workflows/nodes/displayable/bases/api_node/node.py

@@ -92,6 +92,10 @@ class BaseAPINode(BaseNode, Generic[StateType]):
             return self._local_execute_api(data, headers, json, normalized_method, url, timeout)
 
     def _local_execute_api(self, data, headers, json, method, url, timeout):
+        headers = headers or {}
+        if "User-Agent" not in headers:
+            client_headers = self._context.vellum_client._client_wrapper.get_headers()
+            headers["User-Agent"] = client_headers.get("User-Agent")
         try:
             if data is not None:
                 prepped = Request(method=method, url=url, data=data, headers=headers).prepare()
@@ -120,6 +124,11 @@ class BaseAPINode(BaseNode, Generic[StateType]):
     def _vellum_execute_api(self, bearer_token, data, headers, method, url, timeout):
         client_vellum_secret = ClientVellumSecret(name=bearer_token.name) if bearer_token else None
 
+        headers = headers or {}
+        if "User-Agent" not in headers:
+            client_headers = self._context.vellum_client._client_wrapper.get_headers()
+            headers["User-Agent"] = client_headers.get("User-Agent")
+
         # Create request_options if timeout is specified
         request_options = None
         if timeout is not None:

vellum_ai-1.0.8/src/vellum/workflows/nodes/displayable/bases/api_node/tests/test_node.py

@@ -0,0 +1,124 @@
+import pytest
+
+from vellum.client.types.execute_api_response import ExecuteApiResponse
+from vellum.workflows.constants import APIRequestMethod
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.displayable.bases.api_node.node import BaseAPINode
+from vellum.workflows.types.core import VellumSecret
+
+
+@pytest.mark.parametrize("method_value", ["GET", "get", APIRequestMethod.GET])
+def test_api_node_with_string_method(method_value, vellum_client):
+    class TestAPINode(BaseAPINode):
+        method = method_value
+        url = "https://example.com"
+        headers = {"Authorization": VellumSecret(name="API_KEY")}
+
+    mock_response = ExecuteApiResponse(
+        json_={"status": "success"},
+        headers={"content-type": "application/json"},
+        status_code=200,
+        text='{"status": "success"}',
+    )
+    vellum_client.execute_api.return_value = mock_response
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert result.status_code == 200
+
+    vellum_client.execute_api.assert_called_once()
+    call_args = vellum_client.execute_api.call_args
+    assert call_args[1]["method"] == "GET"
+
+
+def test_api_node_with_invalid_method():
+    class TestAPINode(BaseAPINode):
+        method = "INVALID_METHOD"
+        url = "https://example.com"
+
+    node = TestAPINode()
+
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "Invalid HTTP method 'INVALID_METHOD'" == str(exc_info.value)
+
+
+def test_api_node_adds_user_agent_header_when_none_provided(requests_mock):
+    """
+    Tests that the API node adds User-Agent header when no headers are provided.
+    """
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.GET
+        url = "https://example.com/test"
+
+    response_mock = requests_mock.get(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert response_mock.last_request
+    assert "vellum-ai" in response_mock.last_request.headers.get("User-Agent", "")
+
+    assert result.status_code == 200
+
+
+def test_api_node_adds_user_agent_header_when_headers_provided_without_user_agent(requests_mock):
+    """
+    Tests that the API node adds User-Agent header when headers are provided but don't include User-Agent.
+    """
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.POST
+        url = "https://example.com/test"
+        headers = {"Content-Type": "application/json", "Custom-Header": "value"}
+        json = {"test": "data"}
+
+    response_mock = requests_mock.post(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert response_mock.last_request
+    assert "vellum-ai" in response_mock.last_request.headers.get("User-Agent", "")
+    assert response_mock.last_request.headers.get("Content-Type") == "application/json"
+    assert response_mock.last_request.headers.get("Custom-Header") == "value"
+
+    assert result.status_code == 200
+
+
+def test_api_node_preserves_custom_user_agent_header(requests_mock):
+    """
+    Tests that the API node preserves a custom User-Agent header if provided.
+    """
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.GET
+        url = "https://example.com/test"
+        headers = {"User-Agent": "Custom-Agent/1.0"}
+
+    response_mock = requests_mock.get(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert response_mock.last_request
+    assert response_mock.last_request.headers.get("User-Agent") == "Custom-Agent/1.0"
+
+    assert result.status_code == 200

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/src/vellum/workflows/nodes/displayable/tool_calling_node/node.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, Iterator, List, Optional, Set
+from typing import Any, ClassVar, Dict, Iterator, List, Optional, Set, Union
 
 from vellum import ChatMessage, PromptBlock
 from vellum.client.types.prompt_parameters import PromptParameters
@@ -36,7 +36,7 @@ class ToolCallingNode(BaseNode):
     """
 
     ml_model: ClassVar[str] = "gpt-4o-mini"
-    blocks: ClassVar[List[PromptBlock]] = []
+    blocks: ClassVar[List[Union[PromptBlock, Dict[str, Any]]]] = []
     functions: ClassVar[List[Tool]] = []
     prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
     parameters: PromptParameters = DEFAULT_PROMPT_PARAMETERS

vellum_ai-1.0.8/src/vellum/workflows/nodes/displayable/tool_calling_node/tests/test_utils.py

@@ -0,0 +1,236 @@
+import pytest
+from uuid import uuid4
+
+from vellum.client.types.chat_message_prompt_block import ChatMessagePromptBlock
+from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
+from vellum.client.types.initiated_execute_prompt_event import InitiatedExecutePromptEvent
+from vellum.client.types.plain_text_prompt_block import PlainTextPromptBlock
+from vellum.client.types.rich_text_prompt_block import RichTextPromptBlock
+from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.client.types.variable_prompt_block import VariablePromptBlock
+from vellum.prompts.constants import DEFAULT_PROMPT_PARAMETERS
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_tool_router_node, get_function_name
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition, MCPServer, MCPToolDefinition
+
+
+def test_get_function_name_callable():
+    """Test callable"""
+
+    def my_function() -> str:
+        return "test"
+
+    function = my_function
+
+    result = get_function_name(function)
+
+    assert result == "my_function"
+
+
+def test_get_function_name_workflow_class():
+    """Test workflow class."""
+
+    class MyWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        class MyNode(BaseNode):
+            class Outputs(BaseOutputs):
+                result: str
+
+            def run(self) -> Outputs:
+                return self.Outputs(result="test")
+
+        graph = MyNode
+
+    workflow_class = MyWorkflow
+
+    result = get_function_name(workflow_class)
+
+    assert result == "my_workflow"
+
+
+def test_get_function_name_subworkflow_deployment():
+    """Test subworkflow deployment."""
+    deployment_config = DeploymentDefinition(deployment="my-test-deployment", release_tag="v1.0.0")
+
+    result = get_function_name(deployment_config)
+
+    assert result == "mytestdeployment"
+
+
+def test_get_function_name_subworkflow_deployment_uuid():
+    """Test subworkflow deployment with UUID."""
+    deployment_config = DeploymentDefinition(deployment="57f09beb-b463-40e0-bf9e-c972e664352f", release_tag="v1.0.0")
+
+    result = get_function_name(deployment_config)
+
+    assert result == "57f09bebb46340e0bf9ec972e664352f"
+
+
+def test_get_function_name_mcp_tool_definition():
+    """Test MCPToolDefinition function name generation."""
+    mcp_tool = MCPToolDefinition(
+        name="create_repository",
+        server=MCPServer(name="github", url="https://api.github.com"),
+        parameters={"repository_name": "string", "description": "string"},
+    )
+
+    result = get_function_name(mcp_tool)
+
+    assert result == "create_repository"
+
+
+@pytest.mark.parametrize(
+    "toolkit,action,description,expected_result",
+    [
+        ("SLACK", "SLACK_SEND_MESSAGE", "Send message to Slack", "slack_send_message"),
+        ("GMAIL", "GMAIL_CREATE_EMAIL_DRAFT", "Create Gmail draft", "gmail_create_email_draft"),
+    ],
+)
+def test_get_function_name_composio_tool_definition_various_toolkits(
+    toolkit: str, action: str, description: str, expected_result: str
+):
+    """Test ComposioToolDefinition function name generation with various toolkits."""
+    composio_tool = ComposioToolDefinition(toolkit=toolkit, action=action, description=description)
+
+    result = get_function_name(composio_tool)
+
+    assert result == expected_result
+
+
+def test_create_tool_router_node_max_prompt_iterations(vellum_adhoc_prompt_client):
+    # GIVEN a tool router node with max_prompt_iterations set to None
+    tool_router_node = create_tool_router_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[],
+        prompt_inputs=None,
+        parameters=DEFAULT_PROMPT_PARAMETERS,
+        max_prompt_iterations=None,
+    )
+
+    def generate_prompt_events(*args, **kwargs):
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=[StringVellumValue(value="test output")],
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN we run the tool router node
+    node_instance = tool_router_node()
+    outputs = list(node_instance.run())
+    assert outputs[0].name == "results"
+    assert outputs[0].value == [StringVellumValue(type="STRING", value="test output")]
+    assert outputs[1].name == "text"
+    assert outputs[1].value == "test output"
+
+
+def test_create_tool_router_node_chat_history_block_dict(vellum_adhoc_prompt_client):
+    # GIVEN a list of blocks with a chat history block
+    blocks = [
+        {
+            "block_type": "CHAT_MESSAGE",
+            "chat_role": "SYSTEM",
+            "blocks": [
+                {
+                    "block_type": "RICH_TEXT",
+                    "blocks": [{"block_type": "PLAIN_TEXT", "cache_config": None, "text": "first message"}],
+                }
+            ],
+        },
+        {
+            "block_type": "CHAT_MESSAGE",
+            "chat_role": "USER",
+            "blocks": [
+                {
+                    "block_type": "RICH_TEXT",
+                    "blocks": [
+                        {"block_type": "PLAIN_TEXT", "text": "second message"},
+                        {"block_type": "PLAIN_TEXT", "text": "third message"},
+                    ],
+                }
+            ],
+        },
+    ]
+
+    tool_router_node = create_tool_router_node(
+        ml_model="gpt-4o-mini",
+        blocks=blocks,  # type: ignore
+        functions=[],
+        prompt_inputs=None,
+        parameters=DEFAULT_PROMPT_PARAMETERS,
+    )
+
+    def generate_prompt_events(*args, **kwargs):
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=[StringVellumValue(value="test output")],
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN we run the tool router node
+    node_instance = tool_router_node()
+    list(node_instance.run())
+
+    # THEN the API was called with compiled blocks
+    blocks = vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.call_args[1]["blocks"]
+    assert blocks == [
+        ChatMessagePromptBlock(
+            block_type="CHAT_MESSAGE",
+            state=None,
+            cache_config=None,
+            chat_role="SYSTEM",
+            chat_source=None,
+            chat_message_unterminated=None,
+            blocks=[
+                RichTextPromptBlock(
+                    block_type="RICH_TEXT",
+                    state=None,
+                    cache_config=None,
+                    blocks=[
+                        PlainTextPromptBlock(
+                            block_type="PLAIN_TEXT", state=None, cache_config=None, text="first message"
+                        )
+                    ],
+                )
+            ],
+        ),
+        ChatMessagePromptBlock(
+            block_type="CHAT_MESSAGE",
+            state=None,
+            cache_config=None,
+            chat_role="USER",
+            chat_source=None,
+            chat_message_unterminated=None,
+            blocks=[
+                RichTextPromptBlock(
+                    block_type="RICH_TEXT",
+                    state=None,
+                    cache_config=None,
+                    blocks=[
+                        PlainTextPromptBlock(
+                            block_type="PLAIN_TEXT", state=None, cache_config=None, text="second message"
+                        ),
+                        PlainTextPromptBlock(
+                            block_type="PLAIN_TEXT", state=None, cache_config=None, text="third message"
+                        ),
+                    ],
+                )
+            ],
+        ),
+        VariablePromptBlock(block_type="VARIABLE", state=None, cache_config=None, input_variable="chat_history"),
+    ]

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/src/vellum/workflows/nodes/displayable/tool_calling_node/utils.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import Any, Callable, Iterator, List, Optional, Type, Union, cast
+from typing import Any, Callable, Dict, Iterator, List, Optional, Type, Union, cast
 
 from pydash import snake_case
 
@@ -28,7 +28,7 @@ from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state import BaseState
 from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior, Tool
-from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition, MCPToolDefinition
 from vellum.workflows.types.generics import is_workflow_class
 
 CHAT_HISTORY_VARIABLE = "chat_history"
@@ -67,7 +67,7 @@ class ToolRouterNode(InlinePromptNode[ToolCallingState]):
         merge_behavior = MergeBehavior.AWAIT_ATTRIBUTES
 
     def run(self) -> Iterator[BaseOutput]:
-        if self.state.prompt_iterations >= self.max_prompt_iterations:
+        if self.max_prompt_iterations is not None and self.state.prompt_iterations >= self.max_prompt_iterations:
             max_iterations_message = f"Maximum number of prompt iterations `{self.max_prompt_iterations}` reached."
             raise NodeException(message=max_iterations_message, code=WorkflowErrorCode.NODE_EXECUTION)
 
@@ -235,7 +235,7 @@ def _hydrate_composio_tool_definition(tool_def: ComposioToolDefinition) -> Compo
 
 def create_tool_router_node(
     ml_model: str,
-    blocks: List[PromptBlock],
+    blocks: List[Union[PromptBlock, Dict[str, Any]]],
     functions: List[Tool],
     prompt_inputs: Optional[EntityInputsInterface],
     parameters: PromptParameters,
@@ -287,7 +287,16 @@ def create_tool_router_node(
 
     # Add a chat history block to blocks only if one doesn't already exist
     has_chat_history_block = any(
-        block.block_type == "VARIABLE" and block.input_variable == CHAT_HISTORY_VARIABLE for block in blocks
+        (
+            (block["block_type"] if isinstance(block, dict) else block.block_type) == "VARIABLE"
+            and (
+                block["input_variable"]
+                if isinstance(block, dict)
+                else block.input_variable if isinstance(block, VariablePromptBlock) else None
+            )
+            == CHAT_HISTORY_VARIABLE
+        )
+        for block in blocks
     )
 
     if not has_chat_history_block:
@@ -370,6 +379,8 @@ def create_function_node(
             },
         )
         return node
+    elif isinstance(function, MCPToolDefinition):
+        pass
     elif is_workflow_class(function):
         node = type(
             f"DynamicInlineSubworkflowNode_{function.__name__}",
@@ -401,5 +412,7 @@ def get_function_name(function: Tool) -> str:
         return name.replace("-", "")
     elif isinstance(function, ComposioToolDefinition):
         return function.name
+    elif isinstance(function, MCPToolDefinition):
+        return function.name
     else:
         return snake_case(function.__name__)

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/src/vellum/workflows/types/core.py

@@ -13,7 +13,7 @@ from typing import (  # type: ignore[attr-defined]
 )
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
-from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition, MCPToolDefinition
 
 if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
@@ -50,4 +50,4 @@ class ConditionType(Enum):
 
 
 # Type alias for functions that can be called in tool calling nodes
-Tool = Union[Callable[..., Any], DeploymentDefinition, Type["BaseWorkflow"], ComposioToolDefinition]
+Tool = Union[Callable[..., Any], DeploymentDefinition, Type["BaseWorkflow"], ComposioToolDefinition, MCPToolDefinition]

{vellum_ai-1.0.6 → vellum_ai-1.0.8}/src/vellum/workflows/types/definition.py

@@ -8,6 +8,9 @@ from pydantic import BeforeValidator
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.code_resource_definition import CodeResourceDefinition as ClientCodeResourceDefinition
+from vellum.client.types.vellum_secret import VellumSecret
+from vellum.workflows.constants import AuthorizationType
+from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 
 
 def serialize_type_encoder(obj: type) -> Dict[str, Any]:
@@ -118,3 +121,21 @@ class ComposioToolDefinition(UniversalBaseModel):
     def name(self) -> str:
         """Generate a function name for this tool"""
         return self.action.lower()
+
+
+class MCPServer(UniversalBaseModel):
+    name: str
+    url: str
+    authorization_type: AuthorizationType = AuthorizationType.BEARER_TOKEN
+    bearer_token_value: Optional[Union[str, EnvironmentVariableReference]] = None
+    api_key_header_key: Optional[str] = None
+    api_key_header_value: Optional[Union[str, VellumSecret]] = None
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class MCPToolDefinition(UniversalBaseModel):
+    name: str
+    server: MCPServer
+    description: Optional[str] = None
+    parameters: Dict[str, Any] = {}