vellum-ai 0.13.4__py3-none-any.whl → 0.13.6__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.13.4",
+            "X-Fern-SDK-Version": "0.13.6",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/client/types/prompt_node_execution_meta.py

@@ -3,6 +3,7 @@
 from ..core.pydantic_utilities import UniversalBaseModel
 import typing
 from .ml_model_usage import MlModelUsage
+from .price import Price
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
 
@@ -13,6 +14,8 @@ class PromptNodeExecutionMeta(UniversalBaseModel):
     """
 
     usage: typing.Optional[MlModelUsage] = None
+    cost: typing.Optional[Price] = None
+    model_name: typing.Optional[str] = None
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

vellum/client/types/workflow_expand_meta_request.py

@@ -12,6 +12,16 @@ class WorkflowExpandMetaRequest(UniversalBaseModel):
     If enabled, the Prompt Node FULFILLED events will include model host usage tracking. This may increase latency for some model hosts.
     """
 
+    cost: typing.Optional[bool] = pydantic.Field(default=None)
+    """
+    If enabled, the Prompt Node FULFILLED events will include model host cost tracking. This may increase latency for some model hosts.
+    """
+
+    model_name: typing.Optional[bool] = pydantic.Field(default=None)
+    """
+    If enabled, the Prompt Node FULFILLED events will include model host name
+    """
+
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
     else:

{vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/METADATA

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

{vellum_ai-0.13.4.dist-info → vellum_ai-0.13.6.dist-info}/RECORD

@@ -75,22 +75,43 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_n
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py,sha256=1HaEwIbxzWl2OEEVlsNcHFQKMhUKQPlsIMaHnjvMODI,3754
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=5eKEMwrKT5O8dPsOgQ89iZwzpA304jneimQeT9GlpLk,2428
 vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=wQ97k-y9rZwKYlZSA6Z9XGUnGRaJ5BvCFu5QnhuEp4A,7404
-vellum_ee/workflows/display/types.py,sha256=-ninso8Yf2EBn_RkFhpOEOjRhX2bYgCeMWssHoEgPRg,2248
+vellum_ee/workflows/display/types.py,sha256=jWTpfQv3baEWjbP12duc3MBz3TK03Ja1UI5458PVQVY,4211
 vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/utils/vellum.py,sha256=N5VSDYdtHUxvTeSj4zA8aMiAAKn4bAVgpKUDV_obNQ8,3632
 vellum_ee/workflows/display/vellum.py,sha256=8xXRI8b8Tt661H-iZreTQTvLNEKUr4lf-XaKhE7_yUY,8147
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=u5u_KnDp_ktLBlfkT5wiGHIloodSlUT3D3_DQXijtMA,13735
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=cLmW-oMMaP2BSwXCTVJ_KxhKmyfTRRJ_TGhkW3AQYXA,14324
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
 vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=V0edhtohqAWbaHvHkj9Sth4ieaIVejsrsRIr7aCCoVc,16871
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
+vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vellum_ee/workflows/tests/local_files/__init__.py,sha256=UyP6kKkRqr9cTKHQF4MVLdLk5MM9GGcLuqxXsQGm22Y,51
+vellum_ee/workflows/tests/local_files/base_class.py,sha256=UuiC7J68MVr6A4949QYiBpXOLdsvFG_Cw1muEPiHT6I,298
+vellum_ee/workflows/tests/local_files/display/__init__.py,sha256=YISuSFmC8SCTwIRCi08hJS3Us3f95eqkuWUonz26ybk,53
+vellum_ee/workflows/tests/local_files/display/display.py,sha256=GhInaU12K-L_eXFjnRbqnTkvTXr7SQmiBlQbAOOd2PQ,14
+vellum_ee/workflows/tests/local_files/inner_files/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vellum_ee/workflows/tests/local_files/inner_files/inner_class.py,sha256=RUN5k3osh55noeTmvdGK6R6FuHA8zIUX1ME7Zdab0xs,136
+vellum_ee/workflows/tests/local_workflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vellum_ee/workflows/tests/local_workflow/display/__init__.py,sha256=xo75Uqb4aErOsrgy9mJaQwlzLMPDUWIzoi3p-fYk90I,73
+vellum_ee/workflows/tests/local_workflow/display/nodes/__init__.py,sha256=szW_mgOUriyZ6v1vlnevBgkzNi8g83-ihS98UOLHVcE,155
+vellum_ee/workflows/tests/local_workflow/display/nodes/final_output.py,sha256=Kv92TCREiZsB9531KZYaBIq83uHn7e_ECw_yAbD1qfk,1017
+vellum_ee/workflows/tests/local_workflow/display/nodes/templating_node.py,sha256=5cankEe1rDZlXKgILFSPbmN0tUZhIdmcFgz_AguXTJc,1229
+vellum_ee/workflows/tests/local_workflow/display/workflow.py,sha256=JQOxx4xF5yE_xEU7rqbD2i0W8APv16XB6MdsZvlTyZU,2474
+vellum_ee/workflows/tests/local_workflow/inputs.py,sha256=4cgsZBoUbIY0Rs8gknC9yqxQ-sSoULclx_SAm1FT2RA,96
+vellum_ee/workflows/tests/local_workflow/metadata.json,sha256=rdu3h5qkFZiqhCAMxoyoWyMI0O8QALC5-URvSIW6F00,24
+vellum_ee/workflows/tests/local_workflow/nodes/__init__.py,sha256=1F6jxUpSKfPXPj4ZZKSbnX6Mg-VwF3euLJSZfGn6xkM,127
+vellum_ee/workflows/tests/local_workflow/nodes/final_output.py,sha256=ZX7zBv87zirg0w9VKUW3QVDSdBLDqcqAMZjCL_oWbpU,297
+vellum_ee/workflows/tests/local_workflow/nodes/templating_node.py,sha256=NQwFN61QkHfI3Vssz-B0NKGfupK8PU0FDSAIAhYBLi0,325
+vellum_ee/workflows/tests/local_workflow/workflow.py,sha256=A4qOzOPNwePYxWbcAgIPLsmrVS_aVEZEc-wULSv787Q,393
+vellum_ee/workflows/tests/test_display_meta.py,sha256=pzdqND4KLWs7EUIbpXuqgso7BIRpoUsO3T_bgeENs0Q,2205
+vellum_ee/workflows/tests/test_virtual_files.py,sha256=TJEcMR0v2S8CkloXNmCHA0QW0K6pYNGaIjraJz7sFvY,2762
 vellum/__init__.py,sha256=iwoL3PQsiTvtX79J4qlAJ2EIqZ77zYJm3q7o1Ei3Awo,35398
 vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=8nZt88C9SVwWanjLbIQMU3rzb32h5UZfFMBx3VPHB50,111887
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=_dJNXMFx7ZWzJgGQFGTBpowYfKpUBKy4SqXRATAMEOA,1868
+vellum/client/core/client_wrapper.py,sha256=XGwA1YysDKYfgMOJ8jav_tSylN5fJscqdj4e2qXclp4,1868
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 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
 vellum/client/types/prompt_deployment_expand_meta_request.py,sha256=agsiAaHB6lDoZPlnfJ2nmhB4Ud4EiJJTX05YmduyCPo,1910
 vellum/client/types/prompt_deployment_input_request.py,sha256=KrT4-Ew2VvTWXEkYQz2oyHn5EDOgrMW7FzRFaPH3ARg,353
 vellum/client/types/prompt_execution_meta.py,sha256=3hhMZgdAR5mKfnh2e_eVN3oKfT0E9w26khVPrpjn7jk,1141
-vellum/client/types/prompt_node_execution_meta.py,sha256=k8zPFTsXl0iuHi-CNXEQcWP5Adbr8xDv1G1EWZ4KJNM,779
+vellum/client/types/prompt_node_execution_meta.py,sha256=IyWH__nCp5uwS0N32b2ZEsA-Fv7AZDB4nnlRZayU2Gc,888
 vellum/client/types/prompt_node_result.py,sha256=3jewO-nPodoXTq_5RxgwhKfDZrvoPjRZ_vUXLeqiuHY,749
 vellum/client/types/prompt_node_result_data.py,sha256=fNOxBfK3ablDBxkUWVVstJMYaGdHGgu27WxP87E6UQ4,872
 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
 vellum/client/types/workflow_execution_event_type.py,sha256=ESKqV3ItoAlqBooruf-i0AnmEh_GvCySZ0Co3r9Bvt0,170
 vellum/client/types/workflow_execution_node_result_event.py,sha256=wJNwOPwqOl1fBim9Y1rtHiNMoWBZ2xausdSddQI5Mg8,1136
 vellum/client/types/workflow_execution_workflow_result_event.py,sha256=rkv-qaiFEywtb_uSzyClGmq9jSgq-c7-mUznROhCAMU,1147
-vellum/client/types/workflow_expand_meta_request.py,sha256=XOqufKO7gHAumATiJGgxzOYWvGx_iZsf3A49JKs4LJ8,771
+vellum/client/types/workflow_expand_meta_request.py,sha256=-6I1Zveo3wFJEWbmSsN9OiOI7ekCJxF4xEL_ApB6XE8,1151
 vellum/client/types/workflow_node_result_data.py,sha256=fRzWA8CdNCNxk1FuARCJz1rg9LqAwKx-RZYcx5W32Hs,951
 vellum/client/types/workflow_node_result_event.py,sha256=qnhruqkIcKRqP8uB-o-27shcruiT2pjgD7wdrPcp0NE,603
 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
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
 vellum/workflows/workflows/base.py,sha256=k0kUWWko4fHyCqLSU_1cBK_pXZpl9MXekWiG-bdOAo0,18353
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
-vellum_ai-0.13.4.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.13.4.dist-info/METADATA,sha256=B4LvtQbulreSva5kjG54cUi80OaEzT1NMxcVJ4DERP0,5334
-vellum_ai-0.13.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.13.4.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.13.4.dist-info/RECORD,,
+vellum_ai-0.13.6.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.13.6.dist-info/METADATA,sha256=TSCYxziGdaZEtJDxJ0t0JQWbj3UKvtsIeXchhekEhcs,5334
+vellum_ai-0.13.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.13.6.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.13.6.dist-info/RECORD,,

vellum_ee/workflows/display/types.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
+from uuid import UUID
 from typing import TYPE_CHECKING, Dict, Generic, Tuple, Type, TypeVar
 
+from vellum.client.core import UniversalBaseModel
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.ports import Port
@@ -12,6 +14,7 @@ from vellum_ee.workflows.display.base import (
     WorkflowMetaDisplayType,
     WorkflowOutputDisplayType,
 )
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay
 
 if TYPE_CHECKING:
@@ -22,6 +25,18 @@ NodeDisplayType = TypeVar("NodeDisplayType", bound="BaseNodeDisplay")
 WorkflowDisplayType = TypeVar("WorkflowDisplayType", bound="BaseWorkflowDisplay")
 
 
+class NodeDisplay(UniversalBaseModel):
+    input_display: Dict[str, UUID]
+    output_display: Dict[str, UUID]
+    port_display: Dict[str, UUID]
+
+
+class WorkflowEventDisplayContext(UniversalBaseModel):
+    node_displays: Dict[str, NodeDisplay]
+    workflow_inputs: Dict[str, UUID]
+    workflow_outputs: Dict[str, UUID]
+
+
 @dataclass
 class WorkflowDisplayContext(
     Generic[
@@ -48,3 +63,32 @@ class WorkflowDisplayContext(
     workflow_output_displays: Dict[BaseDescriptor, WorkflowOutputDisplayType] = field(default_factory=dict)
     edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplayType] = field(default_factory=dict)
     port_displays: Dict[Port, "PortDisplay"] = field(default_factory=dict)
+
+    def build_event_display_context(self) -> WorkflowEventDisplayContext:
+        workflow_outputs = {
+            output.name: self.workflow_output_displays[output].id for output in self.workflow_output_displays
+        }
+        workflow_inputs = {input.name: self.workflow_input_displays[input].id for input in self.workflow_input_displays}
+        node_displays = {str(node.__id__): self.node_displays[node] for node in self.node_displays}
+        temp_node_displays = {}
+        for node in node_displays:
+            current_node = node_displays[node]
+            input_display = {}
+            if issubclass(current_node.__class__, BaseNodeVellumDisplay):
+                input_display = current_node.node_input_ids_by_name  # type: ignore[attr-defined]
+            node_display_meta = {
+                output.name: current_node.output_display[output].id for output in current_node.output_display
+            }
+            port_display_meta = {port.name: current_node.port_displays[port].id for port in current_node.port_displays}
+
+            temp_node_displays[node] = NodeDisplay(
+                input_display=input_display,
+                output_display=node_display_meta,
+                port_display=port_display_meta,
+            )
+        display_meta = WorkflowEventDisplayContext(
+            workflow_outputs=workflow_outputs,
+            workflow_inputs=workflow_inputs,
+            node_displays=temp_node_displays,
+        )
+        return display_meta

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -1,10 +1,12 @@
 from abc import abstractmethod
 from copy import copy
 from functools import cached_property
+import importlib
 import logging
 from uuid import UUID
-from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, get_args
+from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, Union, get_args
 
+from vellum.workflows import BaseWorkflow
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
@@ -29,7 +31,7 @@ from vellum_ee.workflows.display.base import (
 )
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
-from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
+from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext, WorkflowEventDisplayContext
 
 logger = logging.getLogger(__name__)
 
@@ -333,3 +335,14 @@ class BaseWorkflowDisplay(
 
         workflow_class = get_args(cls.__orig_bases__[0])[0]  # type: ignore [attr-defined]
         cls._workflow_display_registry[workflow_class] = cls
+
+    @staticmethod
+    def gather_event_display_context(
+        module_path: str, workflow_class: Type[BaseWorkflow]
+    ) -> Union[WorkflowEventDisplayContext, None]:
+        workflow_display_module = f"{module_path}.display.workflow"
+        try:
+            display_class = importlib.import_module(workflow_display_module)
+        except ModuleNotFoundError:
+            return None
+        return display_class.WorkflowDisplay(workflow_class).display_context.build_event_display_context()

vellum_ee/workflows/tests/local_files/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa: F401, F403
+
+from .display import *

vellum_ee/workflows/tests/local_files/base_class.py

@@ -0,0 +1,15 @@
+from uuid import uuid4
+
+from .inner_files.inner_class import InnerClass
+
+
+class BaseClass:
+    result = None
+
+    def __init__(self):
+        self.id = uuid4()
+        self.inner_class = InnerClass
+
+    def run(self):
+        self.result = self.inner_class().introduce()
+        return self.result

vellum_ee/workflows/tests/local_files/display/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa: F401, F403
+
+from .display import env

vellum_ee/workflows/tests/local_files/display/display.py

@@ -0,0 +1 @@
+env = "local"

vellum_ee/workflows/tests/local_files/inner_files/inner_class.py

@@ -0,0 +1,9 @@
+import uuid
+
+
+class InnerClass:
+    def __init__(self):
+        self.id = uuid.uuid4()
+
+    def introduce(self):
+        return self.id

vellum_ee/workflows/tests/local_workflow/display/__init__.py

@@ -0,0 +1,4 @@
+# flake8: noqa: F401, F403
+
+from .nodes import *
+from .workflow import *

vellum_ee/workflows/tests/local_workflow/display/nodes/__init__.py

@@ -0,0 +1,4 @@
+from .final_output import FinalOutputDisplay
+from .templating_node import TemplatingNodeDisplay
+
+__all__ = ["TemplatingNodeDisplay", "FinalOutputDisplay"]

vellum_ee/workflows/tests/local_workflow/display/nodes/final_output.py

@@ -0,0 +1,21 @@
+from uuid import UUID
+
+from vellum_ee.workflows.display.nodes import BaseFinalOutputNodeDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
+from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
+
+from ...nodes.final_output import FinalOutput
+
+
+class FinalOutputDisplay(BaseFinalOutputNodeDisplay[FinalOutput]):
+    label = "Final Output"
+    node_id = UUID("f3ef4b2b-fec9-4026-9cc6-e5eac295307f")
+    target_handle_id = UUID("3ec34f6e-da48-40d5-a65b-a48fefa75763")
+    output_id = UUID("5469b810-6ea6-4362-9e79-e360d44a1405")
+    output_name = "final-output"
+    node_input_id = UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")
+    node_input_ids_by_name = {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")}
+    output_display = {
+        FinalOutput.Outputs.value: NodeOutputDisplay(id=UUID("5469b810-6ea6-4362-9e79-e360d44a1405"), name="value")
+    }
+    display_data = NodeDisplayData(position=NodeDisplayPosition(x=2750, y=210), width=459, height=234)

vellum_ee/workflows/tests/local_workflow/display/nodes/templating_node.py

@@ -0,0 +1,27 @@
+from uuid import UUID
+
+from vellum_ee.workflows.display.nodes import BaseTemplatingNodeDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
+
+from ...nodes.templating_node import TemplatingNode
+
+
+class TemplatingNodeDisplay(BaseTemplatingNodeDisplay[TemplatingNode]):
+    label = "Templating Node"
+    node_id = UUID("533c6bd8-6088-4abc-a168-8c1758abcd33")
+    target_handle_id = UUID("048d310e-1a4d-4038-b534-8c68a4509b93")
+    template_input_id = UUID("f97d721a-e685-498e-90c3-9c3d9358fdad")
+    node_input_ids_by_name = {
+        "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
+        "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
+    }
+    output_display = {
+        TemplatingNode.Outputs.result: NodeOutputDisplay(id=UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929"), name="result")
+    }
+    port_displays = {
+        TemplatingNode.Ports.default: PortDisplayOverrides(id=UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62"))
+    }
+    display_data = NodeDisplayData(
+        position=NodeDisplayPosition(x=1970.9393258426965, y=232.4943117977528), width=466, height=224
+    )

vellum_ee/workflows/tests/local_workflow/display/workflow.py

@@ -0,0 +1,58 @@
+from uuid import UUID
+
+from vellum_ee.workflows.display.vellum import (
+    EdgeVellumDisplayOverrides,
+    EntrypointVellumDisplayOverrides,
+    NodeDisplayData,
+    NodeDisplayPosition,
+    WorkflowDisplayData,
+    WorkflowDisplayDataViewport,
+    WorkflowInputsVellumDisplayOverrides,
+    WorkflowMetaVellumDisplayOverrides,
+    WorkflowOutputVellumDisplayOverrides,
+)
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+
+from ..inputs import Inputs
+from ..nodes.final_output import FinalOutput
+from ..nodes.templating_node import TemplatingNode
+from ..workflow import Workflow
+
+
+class WorkflowDisplay(VellumWorkflowDisplay[Workflow]):
+    workflow_display = WorkflowMetaVellumDisplayOverrides(
+        entrypoint_node_id=UUID("0bf86989-13f2-438c-ab9c-d172e5771d31"),
+        entrypoint_node_source_handle_id=UUID("23448f09-81ad-4378-abbd-1cccff350627"),
+        entrypoint_node_display=NodeDisplayData(position=NodeDisplayPosition(x=1545, y=330), width=124, height=48),
+        display_data=WorkflowDisplayData(
+            viewport=WorkflowDisplayDataViewport(x=-1137.2395104895104, y=110.60314685314688, zoom=0.7779720279720279)
+        ),
+    )
+    inputs_display = {
+        Inputs.input_value: WorkflowInputsVellumDisplayOverrides(
+            id=UUID("2268a996-bd17-4832-b3ff-f5662d54b306"), name="input-value", required=True
+        )
+    }
+    entrypoint_displays = {
+        TemplatingNode: EntrypointVellumDisplayOverrides(
+            id=UUID("0bf86989-13f2-438c-ab9c-d172e5771d31"),
+            edge_display=EdgeVellumDisplayOverrides(id=UUID("38532a0e-9432-4ed2-8a34-48a29fd6984d")),
+        )
+    }
+    edge_displays = {
+        (TemplatingNode.Ports.default, FinalOutput): EdgeVellumDisplayOverrides(
+            id=UUID("417c56a4-cdc1-4f9d-a10c-b535163f51e8")
+        )
+    }
+    output_displays = {
+        Workflow.Outputs.final_output: WorkflowOutputVellumDisplayOverrides(
+            id=UUID("5469b810-6ea6-4362-9e79-e360d44a1405"),
+            node_id=UUID("f3ef4b2b-fec9-4026-9cc6-e5eac295307f"),
+            node_input_id=UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb"),
+            name="final-output",
+            label="Final Output",
+            target_handle_id=UUID("3ec34f6e-da48-40d5-a65b-a48fefa75763"),
+            display_data=NodeDisplayData(position=NodeDisplayPosition(x=2750, y=210), width=459, height=234),
+            edge_id=UUID("417c56a4-cdc1-4f9d-a10c-b535163f51e8"),
+        )
+    }

vellum_ee/workflows/tests/local_workflow/inputs.py

@@ -0,0 +1,5 @@
+from vellum.workflows.inputs import BaseInputs
+
+
+class Inputs(BaseInputs):
+    input_value: str

vellum_ee/workflows/tests/local_workflow/metadata.json

@@ -0,0 +1 @@
+{"runner_config": null}

vellum_ee/workflows/tests/local_workflow/nodes/__init__.py

@@ -0,0 +1,4 @@
+from .final_output import FinalOutput
+from .templating_node import TemplatingNode
+
+__all__ = ["TemplatingNode", "FinalOutput"]

vellum_ee/workflows/tests/local_workflow/nodes/final_output.py

@@ -0,0 +1,9 @@
+from vellum.workflows.nodes.displayable import FinalOutputNode
+from vellum.workflows.state import BaseState
+
+from .templating_node import TemplatingNode
+
+
+class FinalOutput(FinalOutputNode[BaseState, str]):
+    class Outputs(FinalOutputNode.Outputs):
+        value = TemplatingNode.Outputs.result

vellum_ee/workflows/tests/local_workflow/nodes/templating_node.py

@@ -0,0 +1,11 @@
+from vellum.workflows.nodes.displayable import TemplatingNode as BaseTemplatingNode
+from vellum.workflows.state import BaseState
+
+from ..inputs import Inputs
+
+
+class TemplatingNode(BaseTemplatingNode[BaseState, str]):
+    template = """{{ example_var_1 }}"""
+    inputs = {
+        "example_var_1": Inputs.input_value,
+    }

vellum_ee/workflows/tests/local_workflow/workflow.py

@@ -0,0 +1,13 @@
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .inputs import Inputs
+from .nodes.final_output import FinalOutput
+from .nodes.templating_node import TemplatingNode
+
+
+class Workflow(BaseWorkflow[Inputs, BaseState]):
+    graph = TemplatingNode >> FinalOutput
+
+    class Outputs(BaseWorkflow.Outputs):
+        final_output = FinalOutput.Outputs.value

vellum_ee/workflows/tests/test_display_meta.py

@@ -0,0 +1,54 @@
+import pytest
+import os
+import sys
+from uuid import UUID, uuid4
+
+from vellum.workflows import BaseWorkflow
+from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
+
+
+@pytest.fixture
+def files() -> dict[str, str]:
+    base_directory = os.path.join(os.path.dirname(__file__), "local_workflow")
+    files = {}
+
+    for root, _, filenames in os.walk(base_directory):
+        for filename in filenames:
+            file_path = os.path.join(root, filename)
+            # Key will be the relative path inside `local_files`
+            relative_path = str(os.path.relpath(file_path, start=base_directory))
+            with open(file_path, encoding="utf-8") as f:
+                files[relative_path] = f.read()
+
+    return files
+
+
+def test_base_class_dynamic_import(files):
+    namespace = str(uuid4())  # Create a unique namespace
+    sys.meta_path.append(VirtualFileFinder(files, namespace))
+
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    display_meta = BaseWorkflowDisplay.gather_event_display_context(namespace, Workflow)
+
+    expected_result = {
+        "node_displays": {
+            "533c6bd8-6088-4abc-a168-8c1758abcd33": {
+                "input_display": {
+                    "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
+                    "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
+                },
+                "output_display": {"result": UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929")},
+                "port_display": {"default": UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62")},
+            },
+            "f3ef4b2b-fec9-4026-9cc6-e5eac295307f": {
+                "input_display": {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")},
+                "output_display": {"value": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+                "port_display": {},
+            },
+        },
+        "workflow_inputs": {"input_value": UUID("2268a996-bd17-4832-b3ff-f5662d54b306")},
+        "workflow_outputs": {"final_output": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+    }
+    assert display_meta
+    assert display_meta.dict() == expected_result

vellum_ee/workflows/tests/test_virtual_files.py

@@ -0,0 +1,82 @@
+import pytest
+import importlib
+import os
+import sys
+from uuid import uuid4
+
+from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
+
+
+@pytest.fixture
+def files() -> dict[str, str]:
+    base_directory = os.path.join(os.path.dirname(__file__), "local_files")
+    files = {}
+
+    for root, _, filenames in os.walk(base_directory):
+        for filename in filenames:
+            file_path = os.path.join(root, filename)
+            # Key will be the relative path inside `local_files`
+            relative_path = str(os.path.relpath(file_path, start=base_directory))
+            with open(file_path, encoding="utf-8") as f:
+                files[relative_path] = f.read()
+
+    return files
+
+
+def test_base_class_dynamic_import(files):
+    """
+    Test dynamically importing the base class and ensuring all required modules are loaded.
+    """
+    namespace = str(uuid4())  # Create a unique namespace
+    sys.meta_path.append(VirtualFileFinder(files, namespace))  # Use virtual file loader
+    base_class_module_path = "base_class.py"
+
+    # When given a valid module path
+    assert base_class_module_path in files, f"{base_class_module_path} is missing from local_files"
+
+    # fully qualified module name for the base class
+    full_module_name = f"{namespace}.base_class"
+
+    # sys.modules keys before import
+    existing_modules = set(sys.modules.keys())
+
+    # If we import the base module (should implicitly load its dependencies)
+    base_class_module = importlib.import_module(full_module_name)
+    assert base_class_module, f"Failed to import module: {full_module_name}"
+
+    baseclass = None
+    for name in dir(base_class_module):
+        if name.startswith("__"):
+            continue
+        obj = getattr(base_class_module, name)
+        if isinstance(obj, type):
+            baseclass = obj
+            break
+
+    assert baseclass
+    # Then we should import all required files
+    new_modules = set(sys.modules.keys()) - existing_modules
+    assert new_modules
+
+    required_modules = [
+        "uuid",
+        f"{namespace}",
+        f"{namespace}.base_class",
+        f"{namespace}.inner_files",
+        f"{namespace}.inner_files.inner_class",
+        f"{namespace}.display",
+        f"{namespace}.display.display",
+    ]
+    for module in required_modules:
+        assert module in sys.modules, f"Module '{module}' was not loaded as expected"
+
+    # Verify that BaseClass can be instantiated without missing imports
+    try:
+        instance = baseclass()
+        assert instance, "Failed to instantiate BaseClass"
+
+        results = instance.run()
+        assert results, "Failed to run BaseClass"
+        assert instance.id
+    except Exception as e:
+        pytest.fail(f"Failed to create an instance of BaseClass: {e}")