vellum-ai 0.13.4__py3-none-any.whl → 0.13.6__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.
Files changed (31) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/client/types/prompt_node_execution_meta.py +3 -0
  3. vellum/client/types/workflow_expand_meta_request.py +10 -0
  4. {vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/METADATA +1 -1
  5. {vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/RECORD +31 -10
  6. vellum_ee/workflows/display/types.py +44 -0
  7. vellum_ee/workflows/display/workflows/base_workflow_display.py +15 -2
  8. vellum_ee/workflows/tests/__init__.py +0 -0
  9. vellum_ee/workflows/tests/local_files/__init__.py +3 -0
  10. vellum_ee/workflows/tests/local_files/base_class.py +15 -0
  11. vellum_ee/workflows/tests/local_files/display/__init__.py +3 -0
  12. vellum_ee/workflows/tests/local_files/display/display.py +1 -0
  13. vellum_ee/workflows/tests/local_files/inner_files/__init__.py +0 -0
  14. vellum_ee/workflows/tests/local_files/inner_files/inner_class.py +9 -0
  15. vellum_ee/workflows/tests/local_workflow/__init__.py +0 -0
  16. vellum_ee/workflows/tests/local_workflow/display/__init__.py +4 -0
  17. vellum_ee/workflows/tests/local_workflow/display/nodes/__init__.py +4 -0
  18. vellum_ee/workflows/tests/local_workflow/display/nodes/final_output.py +21 -0
  19. vellum_ee/workflows/tests/local_workflow/display/nodes/templating_node.py +27 -0
  20. vellum_ee/workflows/tests/local_workflow/display/workflow.py +58 -0
  21. vellum_ee/workflows/tests/local_workflow/inputs.py +5 -0
  22. vellum_ee/workflows/tests/local_workflow/metadata.json +1 -0
  23. vellum_ee/workflows/tests/local_workflow/nodes/__init__.py +4 -0
  24. vellum_ee/workflows/tests/local_workflow/nodes/final_output.py +9 -0
  25. vellum_ee/workflows/tests/local_workflow/nodes/templating_node.py +11 -0
  26. vellum_ee/workflows/tests/local_workflow/workflow.py +13 -0
  27. vellum_ee/workflows/tests/test_display_meta.py +54 -0
  28. vellum_ee/workflows/tests/test_virtual_files.py +82 -0
  29. {vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/LICENSE +0 -0
  30. {vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/WHEEL +0 -0
  31. {vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/entry_points.txt +0 -0
@@ -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.13.4",
21
+ "X-Fern-SDK-Version": "0.13.6",
22
22
  }
23
23
  headers["X_API_KEY"] = self.api_key
24
24
  return headers
@@ -3,6 +3,7 @@
3
3
  from ..core.pydantic_utilities import UniversalBaseModel
4
4
  import typing
5
5
  from .ml_model_usage import MlModelUsage
6
+ from .price import Price
6
7
  from ..core.pydantic_utilities import IS_PYDANTIC_V2
7
8
  import pydantic
8
9
 
@@ -13,6 +14,8 @@ class PromptNodeExecutionMeta(UniversalBaseModel):
13
14
  """
14
15
 
15
16
  usage: typing.Optional[MlModelUsage] = None
17
+ cost: typing.Optional[Price] = None
18
+ model_name: typing.Optional[str] = None
16
19
 
17
20
  if IS_PYDANTIC_V2:
18
21
  model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
@@ -12,6 +12,16 @@ class WorkflowExpandMetaRequest(UniversalBaseModel):
12
12
  If enabled, the Prompt Node FULFILLED events will include model host usage tracking. This may increase latency for some model hosts.
13
13
  """
14
14
 
15
+ cost: typing.Optional[bool] = pydantic.Field(default=None)
16
+ """
17
+ If enabled, the Prompt Node FULFILLED events will include model host cost tracking. This may increase latency for some model hosts.
18
+ """
19
+
20
+ model_name: typing.Optional[bool] = pydantic.Field(default=None)
21
+ """
22
+ If enabled, the Prompt Node FULFILLED events will include model host name
23
+ """
24
+
15
25
  if IS_PYDANTIC_V2:
16
26
  model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
17
27
  else:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-ai
3
- Version: 0.13.4
3
+ Version: 0.13.6
4
4
  Summary:
5
5
  License: MIT
6
6
  Requires-Python: >=3.9,<4.0
@@ -75,22 +75,43 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_n
75
75
  vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py,sha256=1HaEwIbxzWl2OEEVlsNcHFQKMhUKQPlsIMaHnjvMODI,3754
76
76
  vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=5eKEMwrKT5O8dPsOgQ89iZwzpA304jneimQeT9GlpLk,2428
77
77
  vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=wQ97k-y9rZwKYlZSA6Z9XGUnGRaJ5BvCFu5QnhuEp4A,7404
78
- vellum_ee/workflows/display/types.py,sha256=-ninso8Yf2EBn_RkFhpOEOjRhX2bYgCeMWssHoEgPRg,2248
78
+ vellum_ee/workflows/display/types.py,sha256=jWTpfQv3baEWjbP12duc3MBz3TK03Ja1UI5458PVQVY,4211
79
79
  vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
80
80
  vellum_ee/workflows/display/utils/vellum.py,sha256=N5VSDYdtHUxvTeSj4zA8aMiAAKn4bAVgpKUDV_obNQ8,3632
81
81
  vellum_ee/workflows/display/vellum.py,sha256=8xXRI8b8Tt661H-iZreTQTvLNEKUr4lf-XaKhE7_yUY,8147
82
82
  vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
83
- vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=u5u_KnDp_ktLBlfkT5wiGHIloodSlUT3D3_DQXijtMA,13735
83
+ vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=cLmW-oMMaP2BSwXCTVJ_KxhKmyfTRRJ_TGhkW3AQYXA,14324
84
84
  vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
85
85
  vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=V0edhtohqAWbaHvHkj9Sth4ieaIVejsrsRIr7aCCoVc,16871
86
86
  vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
87
87
  vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
88
+ vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
89
+ vellum_ee/workflows/tests/local_files/__init__.py,sha256=UyP6kKkRqr9cTKHQF4MVLdLk5MM9GGcLuqxXsQGm22Y,51
90
+ vellum_ee/workflows/tests/local_files/base_class.py,sha256=UuiC7J68MVr6A4949QYiBpXOLdsvFG_Cw1muEPiHT6I,298
91
+ vellum_ee/workflows/tests/local_files/display/__init__.py,sha256=YISuSFmC8SCTwIRCi08hJS3Us3f95eqkuWUonz26ybk,53
92
+ vellum_ee/workflows/tests/local_files/display/display.py,sha256=GhInaU12K-L_eXFjnRbqnTkvTXr7SQmiBlQbAOOd2PQ,14
93
+ vellum_ee/workflows/tests/local_files/inner_files/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
94
+ vellum_ee/workflows/tests/local_files/inner_files/inner_class.py,sha256=RUN5k3osh55noeTmvdGK6R6FuHA8zIUX1ME7Zdab0xs,136
95
+ vellum_ee/workflows/tests/local_workflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
96
+ vellum_ee/workflows/tests/local_workflow/display/__init__.py,sha256=xo75Uqb4aErOsrgy9mJaQwlzLMPDUWIzoi3p-fYk90I,73
97
+ vellum_ee/workflows/tests/local_workflow/display/nodes/__init__.py,sha256=szW_mgOUriyZ6v1vlnevBgkzNi8g83-ihS98UOLHVcE,155
98
+ vellum_ee/workflows/tests/local_workflow/display/nodes/final_output.py,sha256=Kv92TCREiZsB9531KZYaBIq83uHn7e_ECw_yAbD1qfk,1017
99
+ vellum_ee/workflows/tests/local_workflow/display/nodes/templating_node.py,sha256=5cankEe1rDZlXKgILFSPbmN0tUZhIdmcFgz_AguXTJc,1229
100
+ vellum_ee/workflows/tests/local_workflow/display/workflow.py,sha256=JQOxx4xF5yE_xEU7rqbD2i0W8APv16XB6MdsZvlTyZU,2474
101
+ vellum_ee/workflows/tests/local_workflow/inputs.py,sha256=4cgsZBoUbIY0Rs8gknC9yqxQ-sSoULclx_SAm1FT2RA,96
102
+ vellum_ee/workflows/tests/local_workflow/metadata.json,sha256=rdu3h5qkFZiqhCAMxoyoWyMI0O8QALC5-URvSIW6F00,24
103
+ vellum_ee/workflows/tests/local_workflow/nodes/__init__.py,sha256=1F6jxUpSKfPXPj4ZZKSbnX6Mg-VwF3euLJSZfGn6xkM,127
104
+ vellum_ee/workflows/tests/local_workflow/nodes/final_output.py,sha256=ZX7zBv87zirg0w9VKUW3QVDSdBLDqcqAMZjCL_oWbpU,297
105
+ vellum_ee/workflows/tests/local_workflow/nodes/templating_node.py,sha256=NQwFN61QkHfI3Vssz-B0NKGfupK8PU0FDSAIAhYBLi0,325
106
+ vellum_ee/workflows/tests/local_workflow/workflow.py,sha256=A4qOzOPNwePYxWbcAgIPLsmrVS_aVEZEc-wULSv787Q,393
107
+ vellum_ee/workflows/tests/test_display_meta.py,sha256=pzdqND4KLWs7EUIbpXuqgso7BIRpoUsO3T_bgeENs0Q,2205
108
+ vellum_ee/workflows/tests/test_virtual_files.py,sha256=TJEcMR0v2S8CkloXNmCHA0QW0K6pYNGaIjraJz7sFvY,2762
88
109
  vellum/__init__.py,sha256=iwoL3PQsiTvtX79J4qlAJ2EIqZ77zYJm3q7o1Ei3Awo,35398
89
110
  vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
90
111
  vellum/client/__init__.py,sha256=8nZt88C9SVwWanjLbIQMU3rzb32h5UZfFMBx3VPHB50,111887
91
112
  vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
92
113
  vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
93
- vellum/client/core/client_wrapper.py,sha256=_dJNXMFx7ZWzJgGQFGTBpowYfKpUBKy4SqXRATAMEOA,1868
114
+ vellum/client/core/client_wrapper.py,sha256=XGwA1YysDKYfgMOJ8jav_tSylN5fJscqdj4e2qXclp4,1868
94
115
  vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
95
116
  vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
96
117
  vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -413,7 +434,7 @@ vellum/client/types/prompt_block_state.py,sha256=BRAzTYARoSU36IVZGWMeeqhl5fgFMXC
413
434
  vellum/client/types/prompt_deployment_expand_meta_request.py,sha256=agsiAaHB6lDoZPlnfJ2nmhB4Ud4EiJJTX05YmduyCPo,1910
414
435
  vellum/client/types/prompt_deployment_input_request.py,sha256=KrT4-Ew2VvTWXEkYQz2oyHn5EDOgrMW7FzRFaPH3ARg,353
415
436
  vellum/client/types/prompt_execution_meta.py,sha256=3hhMZgdAR5mKfnh2e_eVN3oKfT0E9w26khVPrpjn7jk,1141
416
- vellum/client/types/prompt_node_execution_meta.py,sha256=k8zPFTsXl0iuHi-CNXEQcWP5Adbr8xDv1G1EWZ4KJNM,779
437
+ vellum/client/types/prompt_node_execution_meta.py,sha256=IyWH__nCp5uwS0N32b2ZEsA-Fv7AZDB4nnlRZayU2Gc,888
417
438
  vellum/client/types/prompt_node_result.py,sha256=3jewO-nPodoXTq_5RxgwhKfDZrvoPjRZ_vUXLeqiuHY,749
418
439
  vellum/client/types/prompt_node_result_data.py,sha256=fNOxBfK3ablDBxkUWVVstJMYaGdHGgu27WxP87E6UQ4,872
419
440
  vellum/client/types/prompt_output.py,sha256=NpDGJNIYIivzQJnBeoJLpJlCk7gqBESLwv5Qtn_20qQ,398
@@ -604,7 +625,7 @@ vellum/client/types/workflow_execution_event_error_code.py,sha256=Vf-MTOx0KclZp3
604
625
  vellum/client/types/workflow_execution_event_type.py,sha256=ESKqV3ItoAlqBooruf-i0AnmEh_GvCySZ0Co3r9Bvt0,170
605
626
  vellum/client/types/workflow_execution_node_result_event.py,sha256=wJNwOPwqOl1fBim9Y1rtHiNMoWBZ2xausdSddQI5Mg8,1136
606
627
  vellum/client/types/workflow_execution_workflow_result_event.py,sha256=rkv-qaiFEywtb_uSzyClGmq9jSgq-c7-mUznROhCAMU,1147
607
- vellum/client/types/workflow_expand_meta_request.py,sha256=XOqufKO7gHAumATiJGgxzOYWvGx_iZsf3A49JKs4LJ8,771
628
+ vellum/client/types/workflow_expand_meta_request.py,sha256=-6I1Zveo3wFJEWbmSsN9OiOI7ekCJxF4xEL_ApB6XE8,1151
608
629
  vellum/client/types/workflow_node_result_data.py,sha256=fRzWA8CdNCNxk1FuARCJz1rg9LqAwKx-RZYcx5W32Hs,951
609
630
  vellum/client/types/workflow_node_result_event.py,sha256=qnhruqkIcKRqP8uB-o-27shcruiT2pjgD7wdrPcp0NE,603
610
631
  vellum/client/types/workflow_node_result_event_state.py,sha256=cC3CdfmXR8bPzSG4W6vDnndA9HAnfF3v4UzmyeVwxog,209
@@ -1419,8 +1440,8 @@ vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528
1419
1440
  vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
1420
1441
  vellum/workflows/workflows/base.py,sha256=k0kUWWko4fHyCqLSU_1cBK_pXZpl9MXekWiG-bdOAo0,18353
1421
1442
  vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
1422
- vellum_ai-0.13.4.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1423
- vellum_ai-0.13.4.dist-info/METADATA,sha256=B4LvtQbulreSva5kjG54cUi80OaEzT1NMxcVJ4DERP0,5334
1424
- vellum_ai-0.13.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1425
- vellum_ai-0.13.4.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1426
- vellum_ai-0.13.4.dist-info/RECORD,,
1443
+ vellum_ai-0.13.6.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1444
+ vellum_ai-0.13.6.dist-info/METADATA,sha256=TSCYxziGdaZEtJDxJ0t0JQWbj3UKvtsIeXchhekEhcs,5334
1445
+ vellum_ai-0.13.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1446
+ vellum_ai-0.13.6.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1447
+ vellum_ai-0.13.6.dist-info/RECORD,,
@@ -1,6 +1,8 @@
1
1
  from dataclasses import dataclass, field
2
+ from uuid import UUID
2
3
  from typing import TYPE_CHECKING, Dict, Generic, Tuple, Type, TypeVar
3
4
 
5
+ from vellum.client.core import UniversalBaseModel
4
6
  from vellum.workflows.descriptors.base import BaseDescriptor
5
7
  from vellum.workflows.nodes import BaseNode
6
8
  from vellum.workflows.ports import Port
@@ -12,6 +14,7 @@ from vellum_ee.workflows.display.base import (
12
14
  WorkflowMetaDisplayType,
13
15
  WorkflowOutputDisplayType,
14
16
  )
17
+ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
15
18
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay
16
19
 
17
20
  if TYPE_CHECKING:
@@ -22,6 +25,18 @@ NodeDisplayType = TypeVar("NodeDisplayType", bound="BaseNodeDisplay")
22
25
  WorkflowDisplayType = TypeVar("WorkflowDisplayType", bound="BaseWorkflowDisplay")
23
26
 
24
27
 
28
+ class NodeDisplay(UniversalBaseModel):
29
+ input_display: Dict[str, UUID]
30
+ output_display: Dict[str, UUID]
31
+ port_display: Dict[str, UUID]
32
+
33
+
34
+ class WorkflowEventDisplayContext(UniversalBaseModel):
35
+ node_displays: Dict[str, NodeDisplay]
36
+ workflow_inputs: Dict[str, UUID]
37
+ workflow_outputs: Dict[str, UUID]
38
+
39
+
25
40
  @dataclass
26
41
  class WorkflowDisplayContext(
27
42
  Generic[
@@ -48,3 +63,32 @@ class WorkflowDisplayContext(
48
63
  workflow_output_displays: Dict[BaseDescriptor, WorkflowOutputDisplayType] = field(default_factory=dict)
49
64
  edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplayType] = field(default_factory=dict)
50
65
  port_displays: Dict[Port, "PortDisplay"] = field(default_factory=dict)
66
+
67
+ def build_event_display_context(self) -> WorkflowEventDisplayContext:
68
+ workflow_outputs = {
69
+ output.name: self.workflow_output_displays[output].id for output in self.workflow_output_displays
70
+ }
71
+ workflow_inputs = {input.name: self.workflow_input_displays[input].id for input in self.workflow_input_displays}
72
+ node_displays = {str(node.__id__): self.node_displays[node] for node in self.node_displays}
73
+ temp_node_displays = {}
74
+ for node in node_displays:
75
+ current_node = node_displays[node]
76
+ input_display = {}
77
+ if issubclass(current_node.__class__, BaseNodeVellumDisplay):
78
+ input_display = current_node.node_input_ids_by_name # type: ignore[attr-defined]
79
+ node_display_meta = {
80
+ output.name: current_node.output_display[output].id for output in current_node.output_display
81
+ }
82
+ port_display_meta = {port.name: current_node.port_displays[port].id for port in current_node.port_displays}
83
+
84
+ temp_node_displays[node] = NodeDisplay(
85
+ input_display=input_display,
86
+ output_display=node_display_meta,
87
+ port_display=port_display_meta,
88
+ )
89
+ display_meta = WorkflowEventDisplayContext(
90
+ workflow_outputs=workflow_outputs,
91
+ workflow_inputs=workflow_inputs,
92
+ node_displays=temp_node_displays,
93
+ )
94
+ return display_meta
@@ -1,10 +1,12 @@
1
1
  from abc import abstractmethod
2
2
  from copy import copy
3
3
  from functools import cached_property
4
+ import importlib
4
5
  import logging
5
6
  from uuid import UUID
6
- from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, get_args
7
+ from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, Union, get_args
7
8
 
9
+ from vellum.workflows import BaseWorkflow
8
10
  from vellum.workflows.descriptors.base import BaseDescriptor
9
11
  from vellum.workflows.edges import Edge
10
12
  from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
@@ -29,7 +31,7 @@ from vellum_ee.workflows.display.base import (
29
31
  )
30
32
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
31
33
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
32
- from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
34
+ from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext, WorkflowEventDisplayContext
33
35
 
34
36
  logger = logging.getLogger(__name__)
35
37
 
@@ -333,3 +335,14 @@ class BaseWorkflowDisplay(
333
335
 
334
336
  workflow_class = get_args(cls.__orig_bases__[0])[0] # type: ignore [attr-defined]
335
337
  cls._workflow_display_registry[workflow_class] = cls
338
+
339
+ @staticmethod
340
+ def gather_event_display_context(
341
+ module_path: str, workflow_class: Type[BaseWorkflow]
342
+ ) -> Union[WorkflowEventDisplayContext, None]:
343
+ workflow_display_module = f"{module_path}.display.workflow"
344
+ try:
345
+ display_class = importlib.import_module(workflow_display_module)
346
+ except ModuleNotFoundError:
347
+ return None
348
+ return display_class.WorkflowDisplay(workflow_class).display_context.build_event_display_context()
File without changes
@@ -0,0 +1,3 @@
1
+ # flake8: noqa: F401, F403
2
+
3
+ from .display import *
@@ -0,0 +1,15 @@
1
+ from uuid import uuid4
2
+
3
+ from .inner_files.inner_class import InnerClass
4
+
5
+
6
+ class BaseClass:
7
+ result = None
8
+
9
+ def __init__(self):
10
+ self.id = uuid4()
11
+ self.inner_class = InnerClass
12
+
13
+ def run(self):
14
+ self.result = self.inner_class().introduce()
15
+ return self.result
@@ -0,0 +1,3 @@
1
+ # flake8: noqa: F401, F403
2
+
3
+ from .display import env
@@ -0,0 +1 @@
1
+ env = "local"
@@ -0,0 +1,9 @@
1
+ import uuid
2
+
3
+
4
+ class InnerClass:
5
+ def __init__(self):
6
+ self.id = uuid.uuid4()
7
+
8
+ def introduce(self):
9
+ return self.id
File without changes
@@ -0,0 +1,4 @@
1
+ # flake8: noqa: F401, F403
2
+
3
+ from .nodes import *
4
+ from .workflow import *
@@ -0,0 +1,4 @@
1
+ from .final_output import FinalOutputDisplay
2
+ from .templating_node import TemplatingNodeDisplay
3
+
4
+ __all__ = ["TemplatingNodeDisplay", "FinalOutputDisplay"]
@@ -0,0 +1,21 @@
1
+ from uuid import UUID
2
+
3
+ from vellum_ee.workflows.display.nodes import BaseFinalOutputNodeDisplay
4
+ from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
5
+ from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
6
+
7
+ from ...nodes.final_output import FinalOutput
8
+
9
+
10
+ class FinalOutputDisplay(BaseFinalOutputNodeDisplay[FinalOutput]):
11
+ label = "Final Output"
12
+ node_id = UUID("f3ef4b2b-fec9-4026-9cc6-e5eac295307f")
13
+ target_handle_id = UUID("3ec34f6e-da48-40d5-a65b-a48fefa75763")
14
+ output_id = UUID("5469b810-6ea6-4362-9e79-e360d44a1405")
15
+ output_name = "final-output"
16
+ node_input_id = UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")
17
+ node_input_ids_by_name = {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")}
18
+ output_display = {
19
+ FinalOutput.Outputs.value: NodeOutputDisplay(id=UUID("5469b810-6ea6-4362-9e79-e360d44a1405"), name="value")
20
+ }
21
+ display_data = NodeDisplayData(position=NodeDisplayPosition(x=2750, y=210), width=459, height=234)
@@ -0,0 +1,27 @@
1
+ from uuid import UUID
2
+
3
+ from vellum_ee.workflows.display.nodes import BaseTemplatingNodeDisplay
4
+ from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplayOverrides
5
+ from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
6
+
7
+ from ...nodes.templating_node import TemplatingNode
8
+
9
+
10
+ class TemplatingNodeDisplay(BaseTemplatingNodeDisplay[TemplatingNode]):
11
+ label = "Templating Node"
12
+ node_id = UUID("533c6bd8-6088-4abc-a168-8c1758abcd33")
13
+ target_handle_id = UUID("048d310e-1a4d-4038-b534-8c68a4509b93")
14
+ template_input_id = UUID("f97d721a-e685-498e-90c3-9c3d9358fdad")
15
+ node_input_ids_by_name = {
16
+ "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
17
+ "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
18
+ }
19
+ output_display = {
20
+ TemplatingNode.Outputs.result: NodeOutputDisplay(id=UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929"), name="result")
21
+ }
22
+ port_displays = {
23
+ TemplatingNode.Ports.default: PortDisplayOverrides(id=UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62"))
24
+ }
25
+ display_data = NodeDisplayData(
26
+ position=NodeDisplayPosition(x=1970.9393258426965, y=232.4943117977528), width=466, height=224
27
+ )
@@ -0,0 +1,58 @@
1
+ from uuid import UUID
2
+
3
+ from vellum_ee.workflows.display.vellum import (
4
+ EdgeVellumDisplayOverrides,
5
+ EntrypointVellumDisplayOverrides,
6
+ NodeDisplayData,
7
+ NodeDisplayPosition,
8
+ WorkflowDisplayData,
9
+ WorkflowDisplayDataViewport,
10
+ WorkflowInputsVellumDisplayOverrides,
11
+ WorkflowMetaVellumDisplayOverrides,
12
+ WorkflowOutputVellumDisplayOverrides,
13
+ )
14
+ from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
15
+
16
+ from ..inputs import Inputs
17
+ from ..nodes.final_output import FinalOutput
18
+ from ..nodes.templating_node import TemplatingNode
19
+ from ..workflow import Workflow
20
+
21
+
22
+ class WorkflowDisplay(VellumWorkflowDisplay[Workflow]):
23
+ workflow_display = WorkflowMetaVellumDisplayOverrides(
24
+ entrypoint_node_id=UUID("0bf86989-13f2-438c-ab9c-d172e5771d31"),
25
+ entrypoint_node_source_handle_id=UUID("23448f09-81ad-4378-abbd-1cccff350627"),
26
+ entrypoint_node_display=NodeDisplayData(position=NodeDisplayPosition(x=1545, y=330), width=124, height=48),
27
+ display_data=WorkflowDisplayData(
28
+ viewport=WorkflowDisplayDataViewport(x=-1137.2395104895104, y=110.60314685314688, zoom=0.7779720279720279)
29
+ ),
30
+ )
31
+ inputs_display = {
32
+ Inputs.input_value: WorkflowInputsVellumDisplayOverrides(
33
+ id=UUID("2268a996-bd17-4832-b3ff-f5662d54b306"), name="input-value", required=True
34
+ )
35
+ }
36
+ entrypoint_displays = {
37
+ TemplatingNode: EntrypointVellumDisplayOverrides(
38
+ id=UUID("0bf86989-13f2-438c-ab9c-d172e5771d31"),
39
+ edge_display=EdgeVellumDisplayOverrides(id=UUID("38532a0e-9432-4ed2-8a34-48a29fd6984d")),
40
+ )
41
+ }
42
+ edge_displays = {
43
+ (TemplatingNode.Ports.default, FinalOutput): EdgeVellumDisplayOverrides(
44
+ id=UUID("417c56a4-cdc1-4f9d-a10c-b535163f51e8")
45
+ )
46
+ }
47
+ output_displays = {
48
+ Workflow.Outputs.final_output: WorkflowOutputVellumDisplayOverrides(
49
+ id=UUID("5469b810-6ea6-4362-9e79-e360d44a1405"),
50
+ node_id=UUID("f3ef4b2b-fec9-4026-9cc6-e5eac295307f"),
51
+ node_input_id=UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb"),
52
+ name="final-output",
53
+ label="Final Output",
54
+ target_handle_id=UUID("3ec34f6e-da48-40d5-a65b-a48fefa75763"),
55
+ display_data=NodeDisplayData(position=NodeDisplayPosition(x=2750, y=210), width=459, height=234),
56
+ edge_id=UUID("417c56a4-cdc1-4f9d-a10c-b535163f51e8"),
57
+ )
58
+ }
@@ -0,0 +1,5 @@
1
+ from vellum.workflows.inputs import BaseInputs
2
+
3
+
4
+ class Inputs(BaseInputs):
5
+ input_value: str
@@ -0,0 +1 @@
1
+ {"runner_config": null}
@@ -0,0 +1,4 @@
1
+ from .final_output import FinalOutput
2
+ from .templating_node import TemplatingNode
3
+
4
+ __all__ = ["TemplatingNode", "FinalOutput"]
@@ -0,0 +1,9 @@
1
+ from vellum.workflows.nodes.displayable import FinalOutputNode
2
+ from vellum.workflows.state import BaseState
3
+
4
+ from .templating_node import TemplatingNode
5
+
6
+
7
+ class FinalOutput(FinalOutputNode[BaseState, str]):
8
+ class Outputs(FinalOutputNode.Outputs):
9
+ value = TemplatingNode.Outputs.result
@@ -0,0 +1,11 @@
1
+ from vellum.workflows.nodes.displayable import TemplatingNode as BaseTemplatingNode
2
+ from vellum.workflows.state import BaseState
3
+
4
+ from ..inputs import Inputs
5
+
6
+
7
+ class TemplatingNode(BaseTemplatingNode[BaseState, str]):
8
+ template = """{{ example_var_1 }}"""
9
+ inputs = {
10
+ "example_var_1": Inputs.input_value,
11
+ }
@@ -0,0 +1,13 @@
1
+ from vellum.workflows import BaseWorkflow
2
+ from vellum.workflows.state import BaseState
3
+
4
+ from .inputs import Inputs
5
+ from .nodes.final_output import FinalOutput
6
+ from .nodes.templating_node import TemplatingNode
7
+
8
+
9
+ class Workflow(BaseWorkflow[Inputs, BaseState]):
10
+ graph = TemplatingNode >> FinalOutput
11
+
12
+ class Outputs(BaseWorkflow.Outputs):
13
+ final_output = FinalOutput.Outputs.value
@@ -0,0 +1,54 @@
1
+ import pytest
2
+ import os
3
+ import sys
4
+ from uuid import UUID, uuid4
5
+
6
+ from vellum.workflows import BaseWorkflow
7
+ from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
8
+ from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
9
+
10
+
11
+ @pytest.fixture
12
+ def files() -> dict[str, str]:
13
+ base_directory = os.path.join(os.path.dirname(__file__), "local_workflow")
14
+ files = {}
15
+
16
+ for root, _, filenames in os.walk(base_directory):
17
+ for filename in filenames:
18
+ file_path = os.path.join(root, filename)
19
+ # Key will be the relative path inside `local_files`
20
+ relative_path = str(os.path.relpath(file_path, start=base_directory))
21
+ with open(file_path, encoding="utf-8") as f:
22
+ files[relative_path] = f.read()
23
+
24
+ return files
25
+
26
+
27
+ def test_base_class_dynamic_import(files):
28
+ namespace = str(uuid4()) # Create a unique namespace
29
+ sys.meta_path.append(VirtualFileFinder(files, namespace))
30
+
31
+ Workflow = BaseWorkflow.load_from_module(namespace)
32
+ display_meta = BaseWorkflowDisplay.gather_event_display_context(namespace, Workflow)
33
+
34
+ expected_result = {
35
+ "node_displays": {
36
+ "533c6bd8-6088-4abc-a168-8c1758abcd33": {
37
+ "input_display": {
38
+ "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
39
+ "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
40
+ },
41
+ "output_display": {"result": UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929")},
42
+ "port_display": {"default": UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62")},
43
+ },
44
+ "f3ef4b2b-fec9-4026-9cc6-e5eac295307f": {
45
+ "input_display": {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")},
46
+ "output_display": {"value": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
47
+ "port_display": {},
48
+ },
49
+ },
50
+ "workflow_inputs": {"input_value": UUID("2268a996-bd17-4832-b3ff-f5662d54b306")},
51
+ "workflow_outputs": {"final_output": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
52
+ }
53
+ assert display_meta
54
+ assert display_meta.dict() == expected_result
@@ -0,0 +1,82 @@
1
+ import pytest
2
+ import importlib
3
+ import os
4
+ import sys
5
+ from uuid import uuid4
6
+
7
+ from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
8
+
9
+
10
+ @pytest.fixture
11
+ def files() -> dict[str, str]:
12
+ base_directory = os.path.join(os.path.dirname(__file__), "local_files")
13
+ files = {}
14
+
15
+ for root, _, filenames in os.walk(base_directory):
16
+ for filename in filenames:
17
+ file_path = os.path.join(root, filename)
18
+ # Key will be the relative path inside `local_files`
19
+ relative_path = str(os.path.relpath(file_path, start=base_directory))
20
+ with open(file_path, encoding="utf-8") as f:
21
+ files[relative_path] = f.read()
22
+
23
+ return files
24
+
25
+
26
+ def test_base_class_dynamic_import(files):
27
+ """
28
+ Test dynamically importing the base class and ensuring all required modules are loaded.
29
+ """
30
+ namespace = str(uuid4()) # Create a unique namespace
31
+ sys.meta_path.append(VirtualFileFinder(files, namespace)) # Use virtual file loader
32
+ base_class_module_path = "base_class.py"
33
+
34
+ # When given a valid module path
35
+ assert base_class_module_path in files, f"{base_class_module_path} is missing from local_files"
36
+
37
+ # fully qualified module name for the base class
38
+ full_module_name = f"{namespace}.base_class"
39
+
40
+ # sys.modules keys before import
41
+ existing_modules = set(sys.modules.keys())
42
+
43
+ # If we import the base module (should implicitly load its dependencies)
44
+ base_class_module = importlib.import_module(full_module_name)
45
+ assert base_class_module, f"Failed to import module: {full_module_name}"
46
+
47
+ baseclass = None
48
+ for name in dir(base_class_module):
49
+ if name.startswith("__"):
50
+ continue
51
+ obj = getattr(base_class_module, name)
52
+ if isinstance(obj, type):
53
+ baseclass = obj
54
+ break
55
+
56
+ assert baseclass
57
+ # Then we should import all required files
58
+ new_modules = set(sys.modules.keys()) - existing_modules
59
+ assert new_modules
60
+
61
+ required_modules = [
62
+ "uuid",
63
+ f"{namespace}",
64
+ f"{namespace}.base_class",
65
+ f"{namespace}.inner_files",
66
+ f"{namespace}.inner_files.inner_class",
67
+ f"{namespace}.display",
68
+ f"{namespace}.display.display",
69
+ ]
70
+ for module in required_modules:
71
+ assert module in sys.modules, f"Module '{module}' was not loaded as expected"
72
+
73
+ # Verify that BaseClass can be instantiated without missing imports
74
+ try:
75
+ instance = baseclass()
76
+ assert instance, "Failed to instantiate BaseClass"
77
+
78
+ results = instance.run()
79
+ assert results, "Failed to run BaseClass"
80
+ assert instance.id
81
+ except Exception as e:
82
+ pytest.fail(f"Failed to create an instance of BaseClass: {e}")