vellum-ai 1.0.11__py3-none-any.whl → 1.1.0__py3-none-any.whl

vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -3,6 +3,7 @@ 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
 
 
@@ -122,3 +123,95 @@ def test_serialize_node__prompt_inputs__mixed_values():
             ],
         },
     }
+
+
+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",
+                            },
+                        },
+                    ],
+                }
+            ],
+        },
+    }

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py

@@ -0,0 +1,98 @@
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_tool_calling_node_mcp_tool.workflow import BasicToolCallingNodeMCPWorkflow
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that uses a generic node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=BasicToolCallingNodeMCPWorkflow)
+
+    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
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+    assert not DeepDiff(
+        [
+            {"id": "3934525a-4807-4dd3-a44f-8569033fece1", "key": "text", "type": "STRING"},
+            {"id": "a6ba35f0-306b-45a4-ad79-c1e9b9a5515c", "key": "chat_history", "type": "CHAT_HISTORY"},
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    tool_calling_node = workflow_raw_data["nodes"][1]
+    function_attributes = next(
+        (attribute for attribute in tool_calling_node["attributes"] if attribute["name"] == "functions"),
+    )
+    assert function_attributes == {
+        "id": "20adf593-c4f0-4c67-8e36-37eb66f28f66",
+        "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": "5a3aef6f-b8a1-4f37-8688-b513da42a35a",
+                            "key": "name",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "github"}},
+                        },
+                        {
+                            "id": "641decf3-79c5-4ef9-9fc8-06570d8a69af",
+                            "key": "url",
+                            "value": {
+                                "type": "CONSTANT_VALUE",
+                                "value": {"type": "STRING", "value": "https://api.githubcopilot.com/mcp/"},
+                            },
+                        },
+                        {
+                            "id": "801a74ca-7966-4ac3-b1b5-bebb71a7de07",
+                            "key": "authorization_type",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "BEARER_TOKEN"}},
+                        },
+                        {
+                            "id": "fcd70e2f-0fb2-4011-a73b-90d5d7643be4",
+                            "key": "bearer_token_value",
+                            "value": {
+                                "type": "ENVIRONMENT_VARIABLE",
+                                "environment_variable": "GITHUB_PERSONAL_ACCESS_TOKEN",
+                            },
+                        },
+                        {
+                            "id": "b2fa2900-0e09-44ff-99db-c5399fd76d28",
+                            "key": "api_key_header_key",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "6ab23414-5f1b-49c1-a0bc-891bbba9124c",
+                            "key": "api_key_header_value",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                    ],
+                }
+            ],
+        },
+    }

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_workflow_deployment_serialization.py

@@ -19,7 +19,7 @@ from tests.workflows.basic_tool_calling_node_workflow_deployment.workflow import
 
 
 def test_serialize_workflow(vellum_client):
-    vellum_client.release_reviews.retrieve_workflow_deployment_release.return_value = WorkflowDeploymentRelease(
+    vellum_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = WorkflowDeploymentRelease(
         id="test-id",
         created=datetime.now(),
         environment=ReleaseEnvironment(

vellum_ee/workflows/display/utils/expressions.py

@@ -358,7 +358,7 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
         }
 
     if isinstance(value, DeploymentDefinition):
-        workflow_deployment_release = display_context.client.release_reviews.retrieve_workflow_deployment_release(
+        workflow_deployment_release = display_context.client.workflow_deployments.retrieve_workflow_deployment_release(
             value.deployment, value.release_tag
         )
         name = workflow_deployment_release.deployment.name or value.deployment

vellum/client/resources/release_reviews/__init__.py

@@ -1,2 +0,0 @@
-# This file was auto-generated by Fern from our API Definition.
-

vellum/client/resources/release_reviews/client.py

@@ -1,139 +0,0 @@
-# This file was auto-generated by Fern from our API Definition.
-
-from ...core.client_wrapper import SyncClientWrapper
-import typing
-from ...core.request_options import RequestOptions
-from ...types.workflow_deployment_release import WorkflowDeploymentRelease
-from ...core.jsonable_encoder import jsonable_encoder
-from ...core.pydantic_utilities import parse_obj_as
-from json.decoder import JSONDecodeError
-from ...core.api_error import ApiError
-from ...core.client_wrapper import AsyncClientWrapper
-
-
-class ReleaseReviewsClient:
-    def __init__(self, *, client_wrapper: SyncClientWrapper):
-        self._client_wrapper = client_wrapper
-
-    def retrieve_workflow_deployment_release(
-        self, id: str, release_id_or_release_tag: str, *, request_options: typing.Optional[RequestOptions] = None
-    ) -> WorkflowDeploymentRelease:
-        """
-        Retrieve a specific Workflow Deployment Release by either its UUID or the name of a Release Tag that points to it.
-
-        Parameters
-        ----------
-        id : str
-            Either the Workflow Deployment's ID or its unique name
-
-        release_id_or_release_tag : str
-            Either the UUID of Workflow Deployment Release you'd like to retrieve, or the name of a Release Tag that's pointing to the Workflow Deployment Release you'd like to retrieve.
-
-        request_options : typing.Optional[RequestOptions]
-            Request-specific configuration.
-
-        Returns
-        -------
-        WorkflowDeploymentRelease
-
-
-        Examples
-        --------
-        from vellum import Vellum
-
-        client = Vellum(
-            api_version="YOUR_API_VERSION",
-            api_key="YOUR_API_KEY",
-        )
-        client.release_reviews.retrieve_workflow_deployment_release(
-            id="id",
-            release_id_or_release_tag="release_id_or_release_tag",
-        )
-        """
-        _response = self._client_wrapper.httpx_client.request(
-            f"v1/workflow-deployments/{jsonable_encoder(id)}/releases/{jsonable_encoder(release_id_or_release_tag)}",
-            base_url=self._client_wrapper.get_environment().default,
-            method="GET",
-            request_options=request_options,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    WorkflowDeploymentRelease,
-                    parse_obj_as(
-                        type_=WorkflowDeploymentRelease,  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            _response_json = _response.json()
-        except JSONDecodeError:
-            raise ApiError(status_code=_response.status_code, body=_response.text)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
-
-
-class AsyncReleaseReviewsClient:
-    def __init__(self, *, client_wrapper: AsyncClientWrapper):
-        self._client_wrapper = client_wrapper
-
-    async def retrieve_workflow_deployment_release(
-        self, id: str, release_id_or_release_tag: str, *, request_options: typing.Optional[RequestOptions] = None
-    ) -> WorkflowDeploymentRelease:
-        """
-        Retrieve a specific Workflow Deployment Release by either its UUID or the name of a Release Tag that points to it.
-
-        Parameters
-        ----------
-        id : str
-            Either the Workflow Deployment's ID or its unique name
-
-        release_id_or_release_tag : str
-            Either the UUID of Workflow Deployment Release you'd like to retrieve, or the name of a Release Tag that's pointing to the Workflow Deployment Release you'd like to retrieve.
-
-        request_options : typing.Optional[RequestOptions]
-            Request-specific configuration.
-
-        Returns
-        -------
-        WorkflowDeploymentRelease
-
-
-        Examples
-        --------
-        import asyncio
-
-        from vellum import AsyncVellum
-
-        client = AsyncVellum(
-            api_version="YOUR_API_VERSION",
-            api_key="YOUR_API_KEY",
-        )
-
-
-        async def main() -> None:
-            await client.release_reviews.retrieve_workflow_deployment_release(
-                id="id",
-                release_id_or_release_tag="release_id_or_release_tag",
-            )
-
-
-        asyncio.run(main())
-        """
-        _response = await self._client_wrapper.httpx_client.request(
-            f"v1/workflow-deployments/{jsonable_encoder(id)}/releases/{jsonable_encoder(release_id_or_release_tag)}",
-            base_url=self._client_wrapper.get_environment().default,
-            method="GET",
-            request_options=request_options,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    WorkflowDeploymentRelease,
-                    parse_obj_as(
-                        type_=WorkflowDeploymentRelease,  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            _response_json = _response.json()
-        except JSONDecodeError:
-            raise ApiError(status_code=_response.status_code, body=_response.text)
-        raise ApiError(status_code=_response.status_code, body=_response_json)

vellum/resources/release_reviews/client.py

@@ -1,3 +0,0 @@
-# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
-
-from vellum.client.resources.release_reviews.client import *