vellum-ai 0.14.47__py3-none-any.whl → 0.14.48__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -18,7 +18,7 @@ class BaseClientWrapper:
18
18
  headers: typing.Dict[str, str] = {
19
19
  "X-Fern-Language": "Python",
20
20
  "X-Fern-SDK-Name": "vellum-ai",
21
- "X-Fern-SDK-Version": "0.14.47",
21
+ "X-Fern-SDK-Version": "0.14.48",
22
22
  }
23
23
  headers["X-API-KEY"] = self.api_key
24
24
  return headers
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-ai
3
- Version: 0.14.47
3
+ Version: 0.14.48
4
4
  Summary:
5
5
  License: MIT
6
6
  Requires-Python: >=3.9,<4.0
@@ -1,9 +1,9 @@
1
1
  vellum_cli/CONTRIBUTING.md,sha256=FtDC7BGxSeMnwCXAUssFsAIElXtmJE-O5Z7BpolcgvI,2935
2
2
  vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
3
- vellum_cli/__init__.py,sha256=7aO9XFnaEVRiVshn86cFudebFUccT-gV8xIARJWqKYo,12257
3
+ vellum_cli/__init__.py,sha256=sCNP_hmFCexEWp1oQdpj8QsIUiAbo4MIlTalZEFruD8,12398
4
4
  vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3363
5
5
  vellum_cli/config.py,sha256=v5BmZ-t_v4Jmqd7KVuQMZF2pRI-rbMspSkVYXIRoTmI,9448
6
- vellum_cli/image_push.py,sha256=8DDvRDJEZ-FukUCqGW1827bg1ybF4xBbx9WyqWYQE-g,6816
6
+ vellum_cli/image_push.py,sha256=skFXf25ixMOX1yfcyAtii-RivYYv-_hsv-Z-bVB6m5Q,7380
7
7
  vellum_cli/init.py,sha256=WpnMXPItPmh0f0bBGIer3p-e5gu8DUGwSArT_FuoMEw,5093
8
8
  vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
9
9
  vellum_cli/ping.py,sha256=p_BCCRjgPhng6JktuECtkDQLbhopt6JpmrtGoLnLJT8,1161
@@ -12,7 +12,7 @@ vellum_cli/push.py,sha256=nWHLDi_w0LXycNkVv00CiNwY469BcTNBn7NphWpCA7E,9711
12
12
  vellum_cli/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
13
13
  vellum_cli/tests/conftest.py,sha256=AFYZryKA2qnUuCPBxBKmHLFoPiE0WhBFFej9tNwSHdc,1526
14
14
  vellum_cli/tests/test_config.py,sha256=uvKGDc8BoVyT9_H0Z-g8469zVxomn6Oi3Zj-vK7O_wU,2631
15
- vellum_cli/tests/test_image_push.py,sha256=i3lJuW8nFRwL1M1OF6752IZYvGAFgKmkB2hd_rYlsmg,2028
15
+ vellum_cli/tests/test_image_push.py,sha256=QM-JlR_aJappvwbCLteQZZf76sd7SE1sRj3armvFK-I,5706
16
16
  vellum_cli/tests/test_init.py,sha256=8UOc_ThfouR4ja5cCl_URuLk7ohr9JXfCnG4yka1OUQ,18754
17
17
  vellum_cli/tests/test_main.py,sha256=qDZG-aQauPwBwM6A2DIu1494n47v3pL28XakTbLGZ-k,272
18
18
  vellum_cli/tests/test_ping.py,sha256=3ucVRThEmTadlV9LrJdCCrr1Ofj3rOjG6ue0BNR2UC0,2523
@@ -26,7 +26,7 @@ vellum_ee/workflows/display/base.py,sha256=EqlQFD56kpqMY02ZBJBQajzJKh33Dwi60Wo77
26
26
  vellum_ee/workflows/display/editor/__init__.py,sha256=MSAgY91xCEg2neH5d8jXx5wRdR962ftZVa6vO9BGq9k,167
27
27
  vellum_ee/workflows/display/editor/types.py,sha256=x-tOOCJ6CF4HmiKDfCmcc3bOVfc1EBlP5o6u5WEfLoY,567
28
28
  vellum_ee/workflows/display/nodes/__init__.py,sha256=jI1aPBQf8DkmrYoZ4O-wR1duqZByOf5mDFmo_wFJPE4,307
29
- vellum_ee/workflows/display/nodes/base_node_display.py,sha256=FGoFpqawVf7mPbMFDZdL9p3qy951Y1RZq8No9NQ8p8Y,15213
29
+ vellum_ee/workflows/display/nodes/base_node_display.py,sha256=1VoNyAw9MzgGePI2wrOShsNccrbS4bTuu3bnTii0Wu4,15480
30
30
  vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=gKAZfc7JBLzcwYPchnpHy2nMVMPmltAszOwLyXDrro0,2085
31
31
  vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
32
32
  vellum_ee/workflows/display/nodes/tests/test_base_node_display.py,sha256=Z4Mf7xLCNiblSbpKI0BrV5modQr-ZcFzhfir_OSyTTs,2997
@@ -40,7 +40,7 @@ vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=MrvyiYD0qgQf
40
40
  vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=m2DmOXm9-jiiIl6zwkXHNfsYp5PTpBHEdt5xaIsabWo,2363
41
41
  vellum_ee/workflows/display/nodes/vellum/final_output_node.py,sha256=jUDI2FwVaw0Or4zJL58J_g0S--i59Hzik60s_Es_M-8,3098
42
42
  vellum_ee/workflows/display/nodes/vellum/guardrail_node.py,sha256=5_5D5PMzBOeUdVtRlANbfEsu7Gv3r37dLvpfjGAqYac,2330
43
- vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=mIwK1gwW7ljrN3dEVCsZlvaGVZhdlbS-hXCMgPf6Ebw,10362
43
+ vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=gVwQwycEPNtCs8tWbFyIMLpCA7zXnqcmuuhFqRWNxZM,10368
44
44
  vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py,sha256=fQV5o83BPTwGX6o-ThN4r7BcIhySyqwpW1JGYWpvSJI,5625
45
45
  vellum_ee/workflows/display/nodes/vellum/map_node.py,sha256=CiklGf5_tDbqE1XQm2mnbtoL01_2JYjcnB4FDTpMImQ,3824
46
46
  vellum_ee/workflows/display/nodes/vellum/merge_node.py,sha256=yBWeN4T_lOsDVnNOKWRiT7JYKu0IR5Fx2z99iq6QKSA,3273
@@ -55,7 +55,7 @@ vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py,sha25
55
55
  vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py,sha256=540FoWMpJ3EN_DPjHsr9ODJWCRVcUa5hZBn-5T2GiHU,1665
56
56
  vellum_ee/workflows/display/nodes/vellum/tests/test_note_node.py,sha256=uiMB0cOxKZzos7YKnj4ef4DFa2bOvZJWIv-hfbUV6Go,1218
57
57
  vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_deployment_node.py,sha256=G-qJyTNJkpqJiEZ3kCJl86CXJINLeFyf2lM0bQHCCOs,3822
58
- vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py,sha256=94ab2dL1adN4ScPxoX3TiYvMguBn0OmwFRq2oCzs7mQ,7468
58
+ vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py,sha256=RPpromm0y9y-MukL8cmxpl9hYaw-JuNo8vFDOcLI4V4,8801
59
59
  vellum_ee/workflows/display/nodes/vellum/tests/test_retry_node.py,sha256=h93ysolmbo2viisyhRnXKHPxiDK0I_dSAbYoHFYIoO4,1953
60
60
  vellum_ee/workflows/display/nodes/vellum/tests/test_subworkflow_deployment_node.py,sha256=BUzHJgjdWnPeZxjFjHfDBKnbFjYjnbXPjc-1hne1B2Y,3965
61
61
  vellum_ee/workflows/display/nodes/vellum/tests/test_templating_node.py,sha256=LSk2gx9TpGXbAqKe8dggQW8yJZqj-Cf0EGJFeGGlEcw,3321
@@ -94,14 +94,14 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_n
94
94
  vellum_ee/workflows/display/types.py,sha256=i4T7ElU5b5h-nA1i3scmEhO1BqmNDc4eJDHavATD88w,2821
95
95
  vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
96
96
  vellum_ee/workflows/display/utils/exceptions.py,sha256=LSwwxCYNxFkf5XMUcFkaZKpQ13OSrI7y_bpEUwbKVk0,169
97
- vellum_ee/workflows/display/utils/expressions.py,sha256=voNX6ZLuLIVu8cCjqASrSnlCZuUL9TVbWkWD6zAMmnE,10644
97
+ vellum_ee/workflows/display/utils/expressions.py,sha256=QKzZpCbVTrDjxO-AKkc-cw7Q2uzYzL4vXpH6OCQAdLk,11731
98
98
  vellum_ee/workflows/display/utils/registry.py,sha256=fWIm5Jj-10gNFjgn34iBu4RWv3Vd15ijtSN0V97bpW8,1513
99
99
  vellum_ee/workflows/display/utils/vellum.py,sha256=HDe1vtIaU35OEg4oC1KC6WNwhw5Laf_mNURorzbeutQ,5418
100
100
  vellum_ee/workflows/display/vellum.py,sha256=o7mq_vk2Yapu9DDKRz5l76h8EmCAypWGQYe6pryrbB8,3576
101
101
  vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
102
102
  vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=NuWlnGNe3Htcfh-l_8e37uitdUsy6WZNB7W7dYcCoUg,33355
103
103
  vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=gxz76AeCqgAZ9D2lZeTiZzxY9eMgn3qOSfVgiqYcOh8,2028
104
- vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=U0LsUqjJ-5qrkqgFsvdWzqWX7OHEMzsDCWrXR1xyzp0,15599
104
+ vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=wkB8tbGdjWyu1PlFfjE0G4y5hw3MvZwfhfN9DE0Ckk4,22142
105
105
  vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=aaKdmWrgEe5YyV4zuDY_4E3y-l59rIHQnNGiPj2OWxQ,359
106
106
  vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
107
107
  vellum_ee/workflows/server/virtual_file_loader.py,sha256=7JphJcSO3H85qiC2DpFfBWjC3JjrbRmoynBC6KKHVsA,2710
@@ -132,7 +132,7 @@ vellum/client/README.md,sha256=qmaVIP42MnxAu8jV7u-CsgVFfs3-pHQODrXdZdFxtaw,4749
132
132
  vellum/client/__init__.py,sha256=PEnFl7LbXQcvAi3bVN2qyt5xm2FtVtq7xWKkcWM3Tg4,120166
133
133
  vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
134
134
  vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
135
- vellum/client/core/client_wrapper.py,sha256=2b05KSUjKpa8HkExXYFatDZrzBB_PJoTnvwdO0cZNvc,1869
135
+ vellum/client/core/client_wrapper.py,sha256=yyzuMHrWz7Em50rdrzRvaaI6oRcCeVgE0RNWHF_GcAo,1869
136
136
  vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
137
137
  vellum/client/core/file.py,sha256=d4NNbX8XvXP32z8KpK2Xovv33nFfruIrpz0QWxlgpZk,2663
138
138
  vellum/client/core/http_client.py,sha256=Z77OIxIbL4OAB2IDqjRq_sYa5yNYAWfmdhdCSSvh6Y4,19552
@@ -1701,8 +1701,8 @@ vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnad
1701
1701
  vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
1702
1702
  vellum/workflows/workflows/tests/test_base_workflow.py,sha256=8P5YIsNMO78_CR1NNK6wkEdkMB4b3Q_Ni1qxh78OnHo,20481
1703
1703
  vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
1704
- vellum_ai-0.14.47.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1705
- vellum_ai-0.14.47.dist-info/METADATA,sha256=ljQBBvIKHzaM01hSSpxXWdIyQPl4GpYpLimvOm1lxCw,5484
1706
- vellum_ai-0.14.47.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1707
- vellum_ai-0.14.47.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1708
- vellum_ai-0.14.47.dist-info/RECORD,,
1704
+ vellum_ai-0.14.48.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1705
+ vellum_ai-0.14.48.dist-info/METADATA,sha256=c8B1dqvOh4tPOrZqgNT8aZdq9fy7L1F010JaK0QD1gE,5484
1706
+ vellum_ai-0.14.48.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1707
+ vellum_ai-0.14.48.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1708
+ vellum_ai-0.14.48.dist-info/RECORD,,
vellum_cli/__init__.py CHANGED
@@ -354,9 +354,10 @@ def images() -> None:
354
354
  help="Tags the provided image inside of Vellum's repo. "
355
355
  "This field does not push multiple local tags of the passed in image.",
356
356
  )
357
- def image_push(image: str, tag: Optional[List[str]] = None) -> None:
357
+ @click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
358
+ def image_push(image: str, tag: Optional[List[str]] = None, workspace: Optional[str] = None) -> None:
358
359
  """Push Docker image to Vellum"""
359
- image_push_command(image, tag)
360
+ image_push_command(image, tag, workspace)
360
361
 
361
362
 
362
363
  @workflows.command(name="init")
vellum_cli/image_push.py CHANGED
@@ -1,5 +1,6 @@
1
1
  import json
2
2
  import logging
3
+ import os
3
4
  import re
4
5
  import subprocess
5
6
  from typing import List, Optional
@@ -9,15 +10,26 @@ from docker import DockerClient
9
10
  from dotenv import load_dotenv
10
11
 
11
12
  from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
13
+ from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
12
14
  from vellum_cli.logger import load_cli_logger
13
15
 
14
16
  _SUPPORTED_ARCHITECTURE = "amd64"
15
17
 
16
18
 
17
- def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
18
- load_dotenv()
19
+ def image_push_command(image: str, tags: Optional[List[str]] = None, workspace: Optional[str] = None) -> None:
20
+ load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
19
21
  logger = load_cli_logger()
20
- vellum_client = create_vellum_client()
22
+ config = load_vellum_cli_config()
23
+ workspace_config = next((w for w in config.workspaces if w.name == workspace), DEFAULT_WORKSPACE_CONFIG)
24
+
25
+ api_key = os.getenv(workspace_config.api_key, None)
26
+ if not api_key:
27
+ raise ValueError(f"API key {workspace_config.api_key} for workspace {workspace} not found")
28
+
29
+ vellum_client = create_vellum_client(
30
+ api_key=api_key,
31
+ api_url=workspace_config.api_url,
32
+ )
21
33
 
22
34
  # Check if we are self hosted by looking at our base url
23
35
  api_url = create_vellum_environment().default
@@ -1,16 +1,37 @@
1
+ import pytest
2
+ import json
3
+ import os
4
+ import shutil
1
5
  import subprocess
6
+ import tempfile
2
7
  from unittest.mock import MagicMock, patch
8
+ from uuid import uuid4
9
+ from typing import Generator
3
10
 
4
11
  from click.testing import CliRunner
12
+ from httpx import Response
5
13
 
6
14
  from vellum_cli import main as cli_main
7
15
 
8
16
 
17
+ @pytest.fixture
18
+ def mock_temp_dir() -> Generator[str, None, None]:
19
+ current_dir = os.getcwd()
20
+ temp_dir = tempfile.mkdtemp()
21
+ os.chdir(temp_dir)
22
+
23
+ yield temp_dir
24
+
25
+ os.chdir(current_dir)
26
+ shutil.rmtree(temp_dir)
27
+
28
+
9
29
  @patch("subprocess.run")
10
30
  @patch("docker.from_env")
11
31
  def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
12
32
  # GIVEN a self hosted vellum api URL env var
13
33
  monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
34
+ monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
14
35
 
15
36
  # Mock Docker client
16
37
  mock_docker_client = MagicMock()
@@ -35,6 +56,94 @@ def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vell
35
56
  assert "Image successfully pushed" in result.output
36
57
 
37
58
 
59
+ @patch("subprocess.run")
60
+ @patch("docker.from_env")
61
+ def test_image_push__self_hosted_happy_path__workspace_option(
62
+ mock_docker_from_env, mock_run, mock_httpx_transport, mock_temp_dir
63
+ ):
64
+ # GIVEN a workspace config with a new env for url
65
+ with open(os.path.join(mock_temp_dir, "vellum.lock.json"), "w") as f:
66
+ f.write(
67
+ json.dumps(
68
+ {
69
+ "workspaces": [
70
+ {
71
+ "name": "my_workspace",
72
+ "api_url": "MY_WORKSPACE_VELLUM_API_URL",
73
+ "api_key": "MY_WORKSPACE_VELLUM_API_KEY",
74
+ }
75
+ ]
76
+ }
77
+ )
78
+ )
79
+
80
+ # AND a .env file with the workspace api key and url
81
+ with open(os.path.join(mock_temp_dir, ".env"), "w") as f:
82
+ f.write(
83
+ "VELLUM_API_KEY=123456abcdef\n"
84
+ "VELLUM_API_URL=https://api.vellum.ai\n"
85
+ "MY_WORKSPACE_VELLUM_API_KEY=789012ghijkl\n"
86
+ "MY_WORKSPACE_VELLUM_API_URL=https://api.vellum.mycompany.ai\n"
87
+ )
88
+
89
+ # AND the Docker client returns the correct response
90
+ mock_docker_client = MagicMock()
91
+ mock_docker_from_env.return_value = mock_docker_client
92
+
93
+ mock_run.side_effect = [
94
+ subprocess.CompletedProcess(
95
+ args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
96
+ ),
97
+ subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
98
+ ]
99
+
100
+ # AND the vellum client returns the correct response for
101
+ mock_httpx_transport.handle_request.side_effect = [
102
+ # First call to get the docker service token
103
+ Response(
104
+ status_code=200,
105
+ text=json.dumps(
106
+ {
107
+ "access_token": "345678mnopqr",
108
+ "organization_id": str(uuid4()),
109
+ "repository": "myrepo.net",
110
+ }
111
+ ),
112
+ ),
113
+ # Second call to push the image
114
+ Response(
115
+ status_code=200,
116
+ text=json.dumps(
117
+ {
118
+ "id": str(uuid4()),
119
+ "name": "myrepo.net/myimage",
120
+ "visibility": "PRIVATE",
121
+ "created": "2021-01-01T00:00:00Z",
122
+ "modified": "2021-01-01T00:00:00Z",
123
+ "repository": "myrepo.net",
124
+ "sha": "sha256:hellosha",
125
+ "tags": [],
126
+ }
127
+ ),
128
+ ),
129
+ ]
130
+
131
+ # WHEN the user runs the image push command
132
+ runner = CliRunner()
133
+ result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest", "--workspace", "my_workspace"])
134
+
135
+ # THEN the command exits successfully
136
+ assert result.exit_code == 0, (result.output, str(result.exception))
137
+
138
+ # AND gives the success message
139
+ assert "Image successfully pushed" in result.output
140
+
141
+ # AND the vellum client was called with the correct api key and url
142
+ request = mock_httpx_transport.handle_request.call_args[0][0]
143
+ assert request.headers["X-API-KEY"] == "789012ghijkl", result.stdout
144
+ assert str(request.url) == "https://api.vellum.mycompany.ai/v1/container-images/push"
145
+
146
+
38
147
  @patch("subprocess.run")
39
148
  @patch("docker.from_env")
40
149
  def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
@@ -260,10 +260,15 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
260
260
 
261
261
  def get_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
262
262
  unadorned_node = get_unadorned_node(self._node)
263
- default_port = unadorned_node.Ports.default
263
+ default_port = next((port for port in unadorned_node.Ports if port.default), None)
264
+ if default_port in port_displays:
265
+ return port_displays[default_port].id
264
266
 
265
- default_port_display = port_displays[default_port]
266
- return default_port_display.id
267
+ first_port = next((port for port in unadorned_node.Ports), None)
268
+ if not first_port:
269
+ raise ValueError(f"Node {self._node.__name__} must have at least one port.")
270
+
271
+ return port_displays[first_port].id
267
272
 
268
273
  def get_trigger_id(self) -> UUID:
269
274
  return self.get_target_handle_id()
@@ -32,7 +32,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
32
32
  output_display = self.output_display[node.Outputs.text]
33
33
  array_display = self.output_display[node.Outputs.results]
34
34
  json_display = self.output_display[node.Outputs.json]
35
- node_blocks = raise_if_descriptor(node.blocks)
35
+ node_blocks = raise_if_descriptor(node.blocks) or []
36
36
  function_definitions = raise_if_descriptor(node.functions)
37
37
 
38
38
  ml_model = str(raise_if_descriptor(node.ml_model))
@@ -6,6 +6,7 @@ from vellum.client.types.variable_prompt_block import VariablePromptBlock
6
6
  from vellum.workflows import BaseWorkflow
7
7
  from vellum.workflows.nodes import BaseNode
8
8
  from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
9
+ from vellum.workflows.ports.port import Port
9
10
  from vellum.workflows.references.lazy import LazyReference
10
11
  from vellum.workflows.state.base import BaseState
11
12
  from vellum_ee.workflows.display.nodes.vellum.inline_prompt_node import BaseInlinePromptNodeDisplay
@@ -215,3 +216,33 @@ def test_serialize_node__unreferenced_variable_block__still_serializes():
215
216
  # warnings = list(workflow_display.errors)
216
217
  # assert len(warnings) == 1
217
218
  # assert "Missing input variable 'foo' for prompt block 0" in str(warnings[0])
219
+
220
+
221
+ def test_serialize_node__port_groups():
222
+ # GIVEN a prompt node with ports
223
+ class MyPromptNode(InlinePromptNode):
224
+ class Ports(InlinePromptNode.Ports):
225
+ apple = Port.on_if(LazyReference(lambda: MyPromptNode.Outputs.text).equals("apple"))
226
+ banana = Port.on_if(LazyReference(lambda: MyPromptNode.Outputs.text).equals("banana"))
227
+
228
+ # AND a workflow with the prompt node
229
+ class MyWorkflow(BaseWorkflow):
230
+ graph = MyPromptNode
231
+
232
+ # WHEN the prompt node is serialized
233
+ workflow_display = get_workflow_display(workflow_class=MyWorkflow)
234
+ serialized_workflow: dict = workflow_display.serialize()
235
+
236
+ # THEN the node should have the ports serialized
237
+ my_prompt_node = next(
238
+ node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
239
+ )
240
+ ports = my_prompt_node["ports"]
241
+ assert len(ports) == 2
242
+ assert ports[0]["id"] == "149d97a4-3da3-44a9-95f7-ea7b8d38b877"
243
+ assert ports[1]["id"] == "71f2d2b3-194f-4492-bc1c-a5ca1f60fb0a"
244
+ assert ports[0]["name"] == "apple"
245
+ assert ports[1]["name"] == "banana"
246
+
247
+ # AND the legacy source_handle_id should be the default port
248
+ assert my_prompt_node["data"]["source_handle_id"] == "149d97a4-3da3-44a9-95f7-ea7b8d38b877"
@@ -1,4 +1,4 @@
1
- from typing import TYPE_CHECKING, Any
1
+ from typing import TYPE_CHECKING, Any, Dict, cast
2
2
 
3
3
  from vellum.client.types.logical_operator import LogicalOperator
4
4
  from vellum.workflows.descriptors.base import BaseDescriptor
@@ -37,7 +37,7 @@ from vellum.workflows.references.output import OutputReference
37
37
  from vellum.workflows.references.state_value import StateValueReference
38
38
  from vellum.workflows.references.vellum_secret import VellumSecretReference
39
39
  from vellum.workflows.references.workflow_input import WorkflowInputReference
40
- from vellum.workflows.types.core import JsonObject
40
+ from vellum.workflows.types.core import JsonArray, JsonObject
41
41
  from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
42
42
 
43
43
  if TYPE_CHECKING:
@@ -235,6 +235,33 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
235
235
  "node_id": str(node_class_display.node_id),
236
236
  }
237
237
 
238
+ if isinstance(value, list):
239
+ serialized_items = [serialize_value(display_context, item) for item in value]
240
+ if all(isinstance(item, dict) and item["type"] == "CONSTANT_VALUE" for item in serialized_items):
241
+ constant_values = []
242
+ for item in serialized_items:
243
+ item_dict = cast(Dict[str, Any], item)
244
+ value_inner = item_dict["value"]
245
+
246
+ if value_inner["type"] == "JSON" and "items" in value_inner:
247
+ # Nested JSON list
248
+ constant_values.append(value_inner["items"])
249
+ else:
250
+ constant_values.append(value_inner["value"])
251
+
252
+ return {
253
+ "type": "CONSTANT_VALUE",
254
+ "value": {
255
+ "type": "JSON",
256
+ "items": constant_values,
257
+ },
258
+ }
259
+ else:
260
+ return {
261
+ "type": "ARRAY_REFERENCE",
262
+ "items": cast(JsonArray, serialized_items), # list[JsonObject] -> JsonArray
263
+ }
264
+
238
265
  if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
239
266
  raise ValueError("Nested references are not supported.")
240
267
 
@@ -425,3 +425,165 @@ def test_serialize_workflow__nested_lazy_reference():
425
425
  "node_id": str(OuterNode.__id__),
426
426
  "node_output_id": str(OuterNode.__output_ids__["bar"]),
427
427
  }
428
+
429
+
430
+ def test_serialize_workflow__array_values():
431
+ # GIVEN a node with array and nested array values
432
+ class MyNode(BaseNode):
433
+ class Outputs(BaseNode.Outputs):
434
+ array_value = ["item1", "item2", "item3"]
435
+ nested_array_value = [["item1", "item2", "item3"], ["item4", "item5", "item6"]]
436
+ mixed_array_value = [["item1"], "item2", "item3"]
437
+
438
+ # AND a workflow that uses these outputs
439
+ class MyWorkflow(BaseWorkflow):
440
+ graph = MyNode
441
+
442
+ class Outputs(BaseWorkflow.Outputs):
443
+ array_output = MyNode.Outputs.array_value
444
+ nested_array_output = MyNode.Outputs.nested_array_value
445
+
446
+ # WHEN we serialize it
447
+ workflow_display = get_workflow_display(workflow_class=MyWorkflow)
448
+ data = workflow_display.serialize()
449
+
450
+ # THEN it should properly serialize the array and dictionary values
451
+ assert isinstance(data["workflow_raw_data"], dict)
452
+ assert isinstance(data["workflow_raw_data"]["nodes"], list)
453
+ raw_nodes = data["workflow_raw_data"]["nodes"]
454
+ generic_nodes = [node for node in raw_nodes if isinstance(node, dict) and node["type"] == "GENERIC"]
455
+ assert len(generic_nodes) > 0
456
+ my_node = generic_nodes[0]
457
+
458
+ outputs = my_node["outputs"]
459
+ assert isinstance(outputs, list)
460
+
461
+ array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "array_value"]
462
+ assert len(array_outputs) > 0
463
+ array_output = array_outputs[0]
464
+
465
+ assert isinstance(array_output, dict)
466
+ assert "value" in array_output
467
+ assert array_output["value"] == {
468
+ "type": "CONSTANT_VALUE",
469
+ "value": {"type": "JSON", "items": ["item1", "item2", "item3"]},
470
+ }
471
+
472
+ nested_array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "nested_array_value"]
473
+ assert len(nested_array_outputs) > 0
474
+ nested_array_output = nested_array_outputs[0]
475
+
476
+ assert isinstance(nested_array_output, dict)
477
+ assert "value" in nested_array_output
478
+ assert nested_array_output["value"] == {
479
+ "type": "CONSTANT_VALUE",
480
+ "value": {"type": "JSON", "items": [["item1", "item2", "item3"], ["item4", "item5", "item6"]]},
481
+ }
482
+
483
+ mixed_array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_array_value"]
484
+ assert len(mixed_array_outputs) > 0
485
+ mixed_array_output = mixed_array_outputs[0]
486
+
487
+ assert isinstance(mixed_array_output, dict)
488
+ assert "value" in mixed_array_output
489
+ assert mixed_array_output["value"] == {
490
+ "type": "CONSTANT_VALUE",
491
+ "value": {"type": "JSON", "items": [["item1"], "item2", "item3"]},
492
+ }
493
+
494
+
495
+ def test_serialize_workflow__array_reference():
496
+ # GIVEN a node with array containing non-constant values (node references)
497
+ class FirstNode(BaseNode):
498
+ class Outputs(BaseNode.Outputs):
499
+ value1: str
500
+ value2: str
501
+
502
+ class SecondNode(BaseNode):
503
+ class Outputs(BaseNode.Outputs):
504
+ # Array containing a mix of constants and node references
505
+ mixed_array = ["constant1", FirstNode.Outputs.value1, "constant2", FirstNode.Outputs.value2]
506
+ mixed_nested_array = [["constant1", FirstNode.Outputs.value1], ["constant2", FirstNode.Outputs.value2]]
507
+
508
+ # AND a workflow that uses these outputs
509
+ class MyWorkflow(BaseWorkflow):
510
+ graph = FirstNode >> SecondNode
511
+
512
+ class Outputs(BaseWorkflow.Outputs):
513
+ mixed_array_output = SecondNode.Outputs.mixed_array
514
+ mixed_nested_array_output = SecondNode.Outputs.mixed_nested_array
515
+
516
+ # WHEN we serialize it
517
+ workflow_display = get_workflow_display(workflow_class=MyWorkflow)
518
+ data = workflow_display.serialize()
519
+
520
+ # THEN it should serialize as an ARRAY_REFERENCE
521
+ assert isinstance(data["workflow_raw_data"], dict)
522
+ assert isinstance(data["workflow_raw_data"]["nodes"], list)
523
+ assert len(data["workflow_raw_data"]["nodes"]) == 5
524
+ second_node = data["workflow_raw_data"]["nodes"][2]
525
+ assert isinstance(second_node, dict)
526
+
527
+ assert "outputs" in second_node
528
+ assert isinstance(second_node["outputs"], list)
529
+ outputs = second_node["outputs"]
530
+
531
+ mixed_array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_array"]
532
+ assert len(mixed_array_outputs) > 0
533
+ mixed_array_output = mixed_array_outputs[0]
534
+
535
+ assert isinstance(mixed_array_output, dict)
536
+ assert "value" in mixed_array_output
537
+ assert mixed_array_output["value"] == {
538
+ "type": "ARRAY_REFERENCE",
539
+ "items": [
540
+ {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant1"}},
541
+ {
542
+ "type": "NODE_OUTPUT",
543
+ "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
544
+ "node_output_id": "419b6afa-fab5-493a-ba1e-4606f4641616",
545
+ },
546
+ {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant2"}},
547
+ {
548
+ "type": "NODE_OUTPUT",
549
+ "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
550
+ "node_output_id": "d1cacc41-478d-49a3-a6b3-1ba2d51291e2",
551
+ },
552
+ ],
553
+ }
554
+
555
+ mixed_nested_array_outputs = [
556
+ val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_nested_array"
557
+ ]
558
+ assert len(mixed_nested_array_outputs) > 0
559
+ mixed_nested_array_output = mixed_nested_array_outputs[0]
560
+
561
+ assert isinstance(mixed_nested_array_output, dict)
562
+ assert "value" in mixed_nested_array_output
563
+ assert mixed_nested_array_output["value"] == {
564
+ "type": "ARRAY_REFERENCE",
565
+ "items": [
566
+ {
567
+ "type": "ARRAY_REFERENCE",
568
+ "items": [
569
+ {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant1"}},
570
+ {
571
+ "type": "NODE_OUTPUT",
572
+ "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
573
+ "node_output_id": "419b6afa-fab5-493a-ba1e-4606f4641616",
574
+ },
575
+ ],
576
+ },
577
+ {
578
+ "type": "ARRAY_REFERENCE",
579
+ "items": [
580
+ {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant2"}},
581
+ {
582
+ "type": "NODE_OUTPUT",
583
+ "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
584
+ "node_output_id": "d1cacc41-478d-49a3-a6b3-1ba2d51291e2",
585
+ },
586
+ ],
587
+ },
588
+ ],
589
+ }