vellum-ai 0.12.1__py3-none-any.whl → 0.12.2__py3-none-any.whl

vellum/client/core/client_wrapper.py

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

vellum/workflows/events/node.py

@@ -10,7 +10,7 @@ from vellum.workflows.ports.port import Port
 from vellum.workflows.references.node import NodeReference
 from vellum.workflows.types.generics import OutputsType
 
-from .types import BaseEvent, default_serializer, serialize_type_encoder
+from .types import BaseEvent, default_serializer, serialize_type_encoder_with_id
 
 if TYPE_CHECKING:
     from vellum.workflows.nodes.bases import BaseNode
@@ -21,7 +21,7 @@ class _BaseNodeExecutionBody(UniversalBaseModel):
 
     @field_serializer("node_definition")
     def serialize_node_definition(self, node_definition: Type, _info: Any) -> Dict[str, Any]:
-        return serialize_type_encoder(node_definition)
+        return serialize_type_encoder_with_id(node_definition)
 
     # Couldn't get this to work with model_config.exclude_none or model_config.exclude_defaults
     # so we're excluding null invoked_ports manually here for now

vellum/workflows/events/tests/test_event.py

@@ -28,6 +28,7 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
 
 
@@ -48,6 +49,8 @@ class MockWorkflow(BaseWorkflow[MockInputs, BaseState]):
 
 name_parts = __name__.split(".")
 module_root = name_parts[: name_parts.index("events")]
+mock_workflow_uuid = str(uuid4_from_hash(MockWorkflow.__qualname__))
+mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
 
 
 @pytest.mark.parametrize(
@@ -73,6 +76,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "workflow.execution.initiated",
                 "body": {
                     "workflow_definition": {
+                        "id": mock_workflow_uuid,
                         "name": "MockWorkflow",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
@@ -113,6 +117,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "node.execution.initiated",
                 "body": {
                     "node_definition": {
+                        "id": mock_node_uuid,
                         "name": "MockNode",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
@@ -122,11 +127,13 @@ module_root = name_parts[: name_parts.index("events")]
                 },
                 "parent": {
                     "node_definition": {
+                        "id": mock_node_uuid,
                         "name": "MockNode",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
                     "parent": {
                         "workflow_definition": {
+                            "id": mock_workflow_uuid,
                             "name": "MockWorkflow",
                             "module": module_root + ["events", "tests", "test_event"],
                         },
@@ -162,6 +169,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "workflow.execution.streaming",
                 "body": {
                     "workflow_definition": {
+                        "id": mock_workflow_uuid,
                         "name": "MockWorkflow",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
@@ -195,6 +203,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "workflow.execution.fulfilled",
                 "body": {
                     "workflow_definition": {
+                        "id": mock_workflow_uuid,
                         "name": "MockWorkflow",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
@@ -228,6 +237,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "workflow.execution.rejected",
                 "body": {
                     "workflow_definition": {
+                        "id": mock_workflow_uuid,
                         "name": "MockWorkflow",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
@@ -262,6 +272,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "node.execution.streaming",
                 "body": {
                     "node_definition": {
+                        "id": mock_node_uuid,
                         "name": "MockNode",
                         "module": module_root + ["events", "tests", "test_event"],
                     },
@@ -296,6 +307,7 @@ module_root = name_parts[: name_parts.index("events")]
                 "name": "node.execution.fulfilled",
                 "body": {
                     "node_definition": {
+                        "id": mock_node_uuid,
                         "name": "MockNode",
                         "module": module_root + ["events", "tests", "test_event"],
                     },

vellum/workflows/events/types.py

@@ -28,6 +28,15 @@ def serialize_type_encoder(obj: type) -> Dict[str, Any]:
     }
 
 
+def serialize_type_encoder_with_id(obj: type) -> Dict[str, Any]:
+    if not hasattr(obj, "__id__"):
+        raise AttributeError(f"The object of type '{type(obj).__name__}' must have an '__id__' attribute.")
+    return {
+        "id": getattr(obj, "__id__"),
+        **serialize_type_encoder(obj),
+    }
+
+
 def default_serializer(obj: Any) -> Any:
     return json.loads(
         json.dumps(
@@ -38,17 +47,18 @@ def default_serializer(obj: Any) -> Any:
 
 
 class CodeResourceDefinition(UniversalBaseModel):
+    id: UUID
     name: str
     module: List[str]
 
     @staticmethod
     def encode(obj: type) -> "CodeResourceDefinition":
-        return CodeResourceDefinition(**serialize_type_encoder(obj))
+        return CodeResourceDefinition(**serialize_type_encoder_with_id(obj))
 
 
 VellumCodeResourceDefinition = Annotated[
     CodeResourceDefinition,
-    BeforeValidator(lambda d: (d if type(d) is dict else serialize_type_encoder(d))),
+    BeforeValidator(lambda d: (d if type(d) is dict else serialize_type_encoder_with_id(d))),
 ]
 
 

vellum/workflows/events/workflow.py

@@ -16,7 +16,7 @@ from .node import (
     NodeExecutionResumedEvent,
     NodeExecutionStreamingEvent,
 )
-from .types import BaseEvent, default_serializer, serialize_type_encoder
+from .types import BaseEvent, default_serializer, serialize_type_encoder_with_id
 
 if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
@@ -27,7 +27,7 @@ class _BaseWorkflowExecutionBody(UniversalBaseModel):
 
     @field_serializer("workflow_definition")
     def serialize_workflow_definition(self, workflow_definition: Type, _info: Any) -> Dict[str, Any]:
-        return serialize_type_encoder(workflow_definition)
+        return serialize_type_encoder_with_id(workflow_definition)
 
 
 class _BaseWorkflowEvent(BaseEvent):

vellum/workflows/nodes/utils.py

@@ -2,11 +2,30 @@ from functools import cache
 from typing import Type
 
 from vellum.workflows.nodes import BaseNode
+from vellum.workflows.ports.port import Port
 from vellum.workflows.types.generics import NodeType
 
 ADORNMENT_MODULE_NAME = "<adornment>"
 
 
+@cache
+def get_unadorned_node(node: Type[BaseNode]) -> Type[BaseNode]:
+    wrapped_node = getattr(node, "__wrapped_node__", None)
+    if wrapped_node is not None:
+        return get_unadorned_node(wrapped_node)
+
+    return node
+
+
+@cache
+def get_unadorned_port(port: Port) -> Port:
+    unadorned_node = get_unadorned_node(port.node_class)
+    if unadorned_node == port.node_class:
+        return port
+
+    return getattr(unadorned_node.Ports, port.name)
+
+
 @cache
 def get_wrapped_node(node: Type[NodeType]) -> Type[BaseNode]:
     wrapped_node = getattr(node, "__wrapped_node__", None)

vellum/workflows/workflows/base.py

@@ -4,6 +4,7 @@ import importlib
 import inspect
 
 from vellum.plugins.utils import load_runtime_plugins
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.event_filters import workflow_event_filter
 
 load_runtime_plugins()
@@ -11,7 +12,7 @@ load_runtime_plugins()
 from datetime import datetime
 from functools import lru_cache
 from threading import Event as ThreadingEvent
-from uuid import uuid4
+from uuid import UUID, uuid4
 from typing import (
     Any,
     Callable,
@@ -85,13 +86,17 @@ class _BaseWorkflowMeta(type):
         if "graph" not in dct:
             dct["graph"] = set()
 
-        return super().__new__(mcs, name, bases, dct)
+        cls = super().__new__(mcs, name, bases, dct)
+        workflow_class = cast(Type["BaseWorkflow"], cls)
+        workflow_class.__id__ = uuid4_from_hash(workflow_class.__qualname__)
+        return workflow_class
 
 
 GraphAttribute = Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]
 
 
 class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkflowMeta):
+    __id__: UUID = uuid4_from_hash(__qualname__)
     graph: ClassVar[GraphAttribute]
     emitters: List[BaseWorkflowEmitter]
     resolvers: List[BaseWorkflowResolver]

vellum/workflows/workflows/event_filters.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Type
 
-from vellum.workflows.events.types import CodeResourceDefinition
+from vellum.workflows.events.types import VellumCodeResourceDefinition
 
 if TYPE_CHECKING:
     from vellum.workflows.events.workflow import WorkflowEvent
@@ -47,7 +47,7 @@ def root_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event:
         return False
 
     event_parent_definition = event.parent.workflow_definition
-    current_workflow_definition = CodeResourceDefinition.encode(workflow_definition)
+    current_workflow_definition = VellumCodeResourceDefinition.encode(workflow_definition)
 
     return event_parent_definition.model_dump() == current_workflow_definition.model_dump()
 

{vellum_ai-0.12.1.dist-info → vellum_ai-0.12.2.dist-info}/METADATA

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

{vellum_ai-0.12.1.dist-info → vellum_ai-0.12.2.dist-info}/RECORD

@@ -20,7 +20,7 @@ vellum_ee/workflows/display/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 vellum_ee/workflows/display/base.py,sha256=3ZFUYRNKL24fBqXhKpa_Dq2W1a-a86J20dmJYA3H2eY,1755
 vellum_ee/workflows/display/nodes/__init__.py,sha256=5XOcZJXYUgaLS55QgRJzyQ_W1tpeprjnYAeYVezqoGw,160
 vellum_ee/workflows/display/nodes/base_node_display.py,sha256=23PLqcpMe_mYkYdug9PDb6Br7o64Thx9-IhcviKGvGo,6662
-vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=E7BFiDv7RdPt_CYhmlgc5ZoQo4JcSQKBGFsAQg0OaJ4,2113
+vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=7F4zQnZQBOTp9turUUS41RaO40Z14m851WOLJHgmHPU,2234
 vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=vyKeJAevAXvEAEtWeTEdBZXD3eJQYW_DEXLKVZ5KmYc,1325
 vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/tests/test_base_node_display.py,sha256=QqR3Ly0RNrXwOeLdW5nERDFt0gRPf76n1bPES6o5UN4,1093
@@ -33,7 +33,7 @@ vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=9JjtXTpyC14X
 vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=ygTjSjYDI4DtkxADWub5rhBnRWItMKWF6fezBrgpOKA,1979
 vellum_ee/workflows/display/nodes/vellum/final_output_node.py,sha256=t5iJQVoRT5g-v2IiUb4kFYdvUVKch0zn27016pzDZoo,2761
 vellum_ee/workflows/display/nodes/vellum/guardrail_node.py,sha256=3TJvHX_Uuf_gr94VkYc_zmNH8I5p71ChIeoAbJZ3ddY,2158
-vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=aU_jZeV0CC3JkIo_LHSh7peLbJEqEaV4odrLPtqzRxI,7375
+vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=E7U4B05-KSlPvTp0nYQTtGPkNa5563VxrNBQmKwvTqw,7267
 vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py,sha256=x5wiuWbRjxNcPGu8BoBEKHwPeqCpHE-vrGjAdM5TJOs,4721
 vellum_ee/workflows/display/nodes/vellum/map_node.py,sha256=AqUlItgSZij12qRKguKVmDbbaLuDy3Cdom5uOlJPqrc,3640
 vellum_ee/workflows/display/nodes/vellum/merge_node.py,sha256=jzO63B9KiEAncnBqmz2ZTcxjmEHozMEe7WnfpcpsQYg,3195
@@ -44,7 +44,7 @@ vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=z
 vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=UNYxoE-89agE8ugK0aWg_uN61jPqlC2VSxWHk568sN4,3324
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=aJfQnIvlrRHVKgQid_gg6VQKkJyPgFnzbvWt9_t0Vz0,3860
-vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=4k9DdjjCkZ_-4qynwg40IipT5Sq2WwgYFFE8WMnMdr0,2612
+vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=hB8dcGMGkuC95kk9hmZUgHsCLwEA37fHTFXj0JzbRjM,4692
 vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=3QsKS5Kht2Nod4A18T0bb4filK9AkyK5oImVlViK3l8,4643
 vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=TEg3QbdE7rLbEhml9pMWmay--phsekGlfGVhTblxCGE,1727
@@ -70,13 +70,13 @@ vellum_ee/workflows/display/vellum.py,sha256=OSv0ZS50h1zJbunJ9TH7VEWFw-exXdK_Zsd
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
 vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=ydAbFMzcY2LURINZbXYm9BAXZdIa3-7rQ86Kupo7qcA,12804
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=AMxNnTm2z3LIR5rqxoCAfuy37F2FTuSRDVtKUoezO8M,1184
-vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=CvxUFudijGW4M7vapm_qrNG4X36PZxzFh6n6qBE9E1Y,17134
+vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=GhIviEMDWNw1p8z20ta08T5PeNCVJs5p2hrOX1uyNxg,17066
 vellum/__init__.py,sha256=QmGeEPXeFxgkZa849KKK3wH3Y641wyt00Rytfay6KiM,35520
 vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=o4m7iRZWEV8rP3GkdaztHAjNmjxjWERlarviFoHzuKI,110927
 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=H2IXjD1b02FWzrMw028nrenbKt081feJ95KpLGbDnFk,1890
+vellum/client/core/client_wrapper.py,sha256=3AlZG0Q4dIs1dgMiXCS_tnLozbk0PAelt6J52g8Tm3k,1890
 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
@@ -1215,11 +1215,11 @@ vellum/workflows/environment/environment.py,sha256=IwcVSwA51jW1JTN_MAN22UbeAtZhi
 vellum/workflows/errors/__init__.py,sha256=tWGPu5xyAU8gRb8_bl0fL7OfU3wxQ9UH6qVwy4X4P_Q,113
 vellum/workflows/errors/types.py,sha256=mvywC2sXtpRrkDRfTx-O6aXAp0hOFyx27a1tqguKt4g,2493
 vellum/workflows/events/__init__.py,sha256=6pxxceJo2dcaRkWtkDAYlUQZ-PHBQSZytIoyuUK48Qw,759
-vellum/workflows/events/node.py,sha256=QOikR6w5KXJUjnDNF2RhjXRbG7lJySetAoPamyr733Y,5264
+vellum/workflows/events/node.py,sha256=uHT6If0esgZ3nLjrjmUPTKf3qbjGhoV_x5YKpjDBDcU,5280
 vellum/workflows/events/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/events/tests/test_event.py,sha256=EFg34nGpl_LvbHWN5AuiMqH6PetwGF_z0TlETGsn_bU,12725
-vellum/workflows/events/types.py,sha256=7sBSMNstsibqq93Klu4A-GInzo4IB23xWel595Uyi_Y,3077
-vellum/workflows/events/workflow.py,sha256=LW1qjLYCa9mpZjxp6YuDDA0M-q0COBIfLuduP9pOfyo,5224
+vellum/workflows/events/tests/test_event.py,sha256=izB6Y9U5ROgmHBBpLNUY2navK4-qFp6hdJqJNz6Beek,13350
+vellum/workflows/events/types.py,sha256=cjRE8WL8tYCFradd9NOGl_H0mN3LiWWnA1uHmyT2Q0Q,3412
+vellum/workflows/events/workflow.py,sha256=l5tXes0sg7iWaA1ZUE5dtAqNnGQ8iy6trVbOU9meu7U,5240
 vellum/workflows/exceptions.py,sha256=gXQvDL919cK3pwXc8BhLGKmR-YOskYDLvL1zPSSmLS4,579
 vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/expressions/accessor.py,sha256=OFvAHAVABr-k7GceIhtzIurV4OuV_yHft7JPRsq87Es,1472
@@ -1327,7 +1327,7 @@ vellum/workflows/nodes/displayable/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5
 vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py,sha256=lLnXKAUYtgvQ6MpT4GoTrqLtdlyDlUt1pPHrmu-Gf00,4705
 vellum/workflows/nodes/displayable/tests/test_search_node_wth_text_output.py,sha256=4CMwDtXwTaEvFfDpA6j2iLqc7S6IICSkvVZOobEpeps,6954
 vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py,sha256=KqKJtJ0vuNoPuUPMdILmBTt4a2fBBxxun-nmOI7T8jo,2585
-vellum/workflows/nodes/utils.py,sha256=RWQbcc6MgxHNwFjtasY4hzIH8pKe0E3zfVt6eJKGTp4,587
+vellum/workflows/nodes/utils.py,sha256=EZt7CzJmgQBR_GWFpZr8d-oaoti3tolTd2Cv9wm7dKo,1087
 vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
 vellum/workflows/outputs/base.py,sha256=a7W6rNSDSawwGAXYjNTF2iHb9lnZu7WFSOagZIyy__k,7976
 vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
@@ -1375,10 +1375,10 @@ vellum/workflows/utils/uuids.py,sha256=DFzPv9RCvsKhvdTEIQyfSek2A31D6S_QcmeLPbgrg
 vellum/workflows/utils/vellum_variables.py,sha256=DsjVj_M_vTafpi5OUDs4KNrmbU2n4LR7fcLhuHb67Z4,3123
 vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528t75s,683
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
-vellum/workflows/workflows/base.py,sha256=jFu_AAyt4IqOrwanyflNnDhgYpX4oTVR5xnow8ynRhg,16811
-vellum/workflows/workflows/event_filters.py,sha256=-uQcMB7IpPd-idMku8f2QNVhPXPFWo6FZLlGjRf8rCo,1996
-vellum_ai-0.12.1.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.12.1.dist-info/METADATA,sha256=m7888sp693GSPQLWUuFa1W1ehpB2orR-RuEqmqpV-SQ,5128
-vellum_ai-0.12.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.12.1.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.12.1.dist-info/RECORD,,
+vellum/workflows/workflows/base.py,sha256=zpspOEdO5Ye_0ZvN-Wkzv9iQSiF1sD201ba8lhbnPbs,17086
+vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
+vellum_ai-0.12.2.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.12.2.dist-info/METADATA,sha256=alKUy1QGixY3flPix2H1o-GW9kh_xrF5CKtOm5EhJ6I,5128
+vellum_ai-0.12.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.12.2.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.12.2.dist-info/RECORD,,

vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -1,6 +1,7 @@
 from uuid import UUID
 from typing import ClassVar, Dict, Optional
 
+from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.ports import Port
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.utils.uuids import uuid4_from_hash
@@ -40,6 +41,8 @@ class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
         return self.get_target_handle_id()
 
     def get_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
-        default_port = self._node.Ports.default
+        unadorned_node = get_unadorned_node(self._node)
+        default_port = unadorned_node.Ports.default
+
         default_port_display = port_displays[default_port]
         return default_port_display.id

vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -1,9 +1,8 @@
 from uuid import UUID
-from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 from vellum import PromptBlock, RichTextChildBlock, VellumVariable
 from vellum.workflows.nodes import InlinePromptNode
-from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
@@ -30,8 +29,8 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
         node_inputs, prompt_inputs = self._generate_node_and_prompt_inputs(node_id, node, display_context)
         input_variable_id_by_name = {prompt_input.key: prompt_input.id for prompt_input in prompt_inputs}
 
-        _, output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.text)]
-        _, array_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.results)]
+        _, output_display = display_context.node_output_displays[node.Outputs.text]
+        _, array_display = display_context.node_output_displays[node.Outputs.results]
         node_blocks = raise_if_descriptor(node.blocks)
 
         return {

vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -1,11 +1,13 @@
 from uuid import UUID
-from typing import Any, ClassVar, Generic, Optional, TypeVar
+from typing import Any, Callable, ClassVar, Generic, Optional, Type, TypeVar, cast
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME, get_wrapped_node
 from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
@@ -54,3 +56,43 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
                 serialized_node_definition["name"] = node.__name__
 
         return serialized_node
+
+    @classmethod
+    def wrap(cls, error_output_id: Optional[UUID] = None) -> Callable[..., Type["BaseTryNodeDisplay"]]:
+        _error_output_id = error_output_id
+
+        NodeDisplayType = TypeVar("NodeDisplayType", bound=BaseNodeDisplay)
+
+        def decorator(inner_cls: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
+            node_class = inner_cls.infer_node_class()
+            wrapped_node_class = cast(Type[BaseNode], node_class.__wrapped_node__)
+
+            # Mypy gets mad about dynamic parameter types like this, but it's fine
+            class TryNodeDisplay(BaseTryNodeDisplay[node_class]):  # type: ignore[valid-type]
+                error_output_id = _error_output_id
+
+            setattr(inner_cls, "__adorned_by__", TryNodeDisplay)
+
+            # We must edit the node display class to use __wrapped_node__ everywhere it
+            # references the adorned node class, which is three places:
+
+            # 1. The node display class' parameterized type
+            original_base_node_display = get_original_base(inner_cls)
+            original_base_node_display.__args__ = (wrapped_node_class,)
+            inner_cls._node_display_registry[wrapped_node_class] = inner_cls
+
+            # 2. The node display class' output displays
+            old_outputs = list(inner_cls.output_display.keys())
+            for old_output in old_outputs:
+                new_output = getattr(wrapped_node_class.Outputs, old_output.name)
+                inner_cls.output_display[new_output] = inner_cls.output_display.pop(old_output)
+
+            # 3. The node display class' port displays
+            old_ports = list(inner_cls.port_displays.keys())
+            for old_port in old_ports:
+                new_port = getattr(wrapped_node_class.Ports, old_port.name)
+                inner_cls.port_displays[new_port] = inner_cls.port_displays.pop(old_port)
+
+            return inner_cls
+
+        return decorator

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -6,7 +6,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node, has_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
@@ -352,17 +352,12 @@ class VellumWorkflowDisplay(
         port_displays: Dict[Port, PortDisplay],
         overrides: Optional[EdgeVellumDisplayOverrides] = None,
     ) -> EdgeVellumDisplay:
-        source_node = edge.from_port.node_class
-        target_node = edge.to_node
-
-        if has_wrapped_node(source_node):
-            source_node = get_wrapped_node(source_node)
-
-        if has_wrapped_node(target_node):
-            target_node = get_wrapped_node(target_node)
+        source_node = get_unadorned_node(edge.from_port.node_class)
+        target_node = get_unadorned_node(edge.to_node)
 
         source_node_id = node_displays[source_node].node_id
-        source_handle_id = port_displays[edge.from_port].id
+        from_port = get_unadorned_port(edge.from_port)
+        source_handle_id = port_displays[from_port].id
 
         target_node_display = node_displays[target_node]
         target_node_id = target_node_display.node_id