vellum-ai 0.11.10__tar.gz → 0.12.0__tar.gz

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/PKG-INFO

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

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -1,18 +1,30 @@
 from functools import cached_property
 import inspect
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, cast, get_args, get_origin
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    ForwardRef,
+    Generic,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    cast,
+    get_args,
+    get_origin,
+)
 
 from vellum.workflows.nodes.bases.base import BaseNode
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.names import pascal_to_title_case
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import CodeResourceDefinition, NodeDefinition
 
 if TYPE_CHECKING:
@@ -21,16 +33,24 @@ if TYPE_CHECKING:
 _NodeDisplayAttrType = TypeVar("_NodeDisplayAttrType")
 
 
-class BaseNodeDisplay(Generic[NodeType]):
+class BaseNodeDisplayMeta(type):
+    def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        cls = super().__new__(mcs, name, bases, dct)
+        if isinstance(dct.get("node_id"), UUID):
+            # Display classes are able to override the id of the node class it's parameterized by
+            base_node_display_class = cast(Type["BaseNodeDisplay"], cls)
+            node_class = base_node_display_class.infer_node_class()
+            node_class.__id__ = dct["node_id"]
+        return cls
+
+
+class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     output_display: Dict[OutputReference, NodeOutputDisplay] = {}
     port_displays: Dict[Port, PortDisplayOverrides] = {}
 
     # Used to store the mapping between node types and their display classes
     _node_display_registry: Dict[Type[NodeType], Type["BaseNodeDisplay"]] = {}
 
-    def __init__(self, node: Type[NodeType]):
-        self._node = node
-
     def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
         raise NotImplementedError(f"Serialization for nodes of type {self._node.__name__} is not supported.")
 
@@ -71,6 +91,23 @@ class BaseNodeDisplay(Generic[NodeType]):
     def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Type["BaseNodeDisplay"]:
         return cls._node_display_registry[node_class]
 
+    @classmethod
+    def infer_node_class(cls) -> Type[NodeType]:
+        original_base = get_original_base(cls)
+        node_class = get_args(original_base)[0]
+        if isinstance(node_class, TypeVar):
+            bounded_class = node_class.__bound__
+            if inspect.isclass(bounded_class) and issubclass(bounded_class, BaseNode):
+                return cast(Type[NodeType], bounded_class)
+
+            if isinstance(bounded_class, ForwardRef) and bounded_class.__forward_arg__ == BaseNode.__name__:
+                return cast(Type[NodeType], BaseNode)
+
+        if issubclass(node_class, BaseNode):
+            return node_class
+
+        raise ValueError(f"Node {cls.__name__} must be a subclass of {BaseNode.__name__}")
+
     @cached_property
     def node_id(self) -> UUID:
         """Can be overridden as a class attribute to specify a custom node id."""
@@ -81,6 +118,10 @@ class BaseNodeDisplay(Generic[NodeType]):
         """Can be overridden as a class attribute to specify a custom label."""
         return pascal_to_title_case(self._node.__name__)
 
+    @property
+    def _node(self) -> Type[NodeType]:
+        return self.infer_node_class()
+
     @classmethod
     def _get_explicit_node_display_attr(
         cls,
@@ -124,17 +165,5 @@ class BaseNodeDisplay(Generic[NodeType]):
     def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)
 
-        original_base = get_original_base(cls)
-        node_class = get_args(original_base)[0]
-        if isinstance(node_class, TypeVar):
-            bounded_class = node_class.__bound__
-            if inspect.isclass(bounded_class) and issubclass(bounded_class, BaseNode):
-                cls._node_display_registry[bounded_class] = cls
-        elif issubclass(node_class, BaseNode):
-            if has_wrapped_node(node_class):
-                wrapped_node = get_wrapped_node(node_class)
-                if wrapped_node._is_wrapped_node:
-                    cls._node_display_registry[wrapped_node] = cls
-                    return
-
-            cls._node_display_registry[node_class] = cls
+        node_class = cls.infer_node_class()
+        cls._node_display_registry[node_class] = cls

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -3,9 +3,9 @@ from typing import ClassVar, Dict, Optional
 
 from vellum.workflows.ports import Port
 from vellum.workflows.types.generics import NodeType
+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.types import PortDisplay
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import NodeDisplayData
 
 

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,3 +1,4 @@
+import types
 from typing import Optional, Type
 
 from vellum.workflows.types.generics import NodeType
@@ -11,9 +12,17 @@ def get_node_display_class(
         node_display_class = base_class.get_from_node_display_registry(node_class)
     except KeyError:
         try:
-            return get_node_display_class(
+            base_node_display_class = get_node_display_class(
                 base_class, node_class.__bases__[0], node_class if root_node_class is None else root_node_class
             )
+
+            # `base_node_display_class` is always a Generic class, so it's safe to index into it
+            NodeDisplayBaseClass = base_node_display_class[node_class]  # type: ignore[index]
+            NodeDisplayClass = types.new_class(
+                f"{node_class.__name__}Display",
+                bases=(NodeDisplayBaseClass,),
+            )
+            return NodeDisplayClass
         except IndexError:
             return base_class
 

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -15,7 +15,7 @@ def node_with_implicit_properties():
 
     expected_id = UUID("ace7f746-4fe6-45c7-8207-fc8a4d0c7f6f")
 
-    return MyNode, MyNodeDisplay, expected_id
+    return MyNodeDisplay, expected_id
 
 
 @pytest.fixture
@@ -28,7 +28,7 @@ def node_with_explicit_properties():
     class MyNodeDisplay(BaseNodeDisplay[MyNode]):
         node_id = explicit_id
 
-    return MyNode, MyNodeDisplay, explicit_id
+    return MyNodeDisplay, explicit_id
 
 
 @pytest.fixture(
@@ -42,6 +42,7 @@ def node_info(request):
 
 
 def test_get_id(node_info):
-    node, node_display, expected_id = node_info
+    node_display, expected_id = node_info
 
-    assert node_display(node).node_id == expected_id
+    assert node_display().node_id == expected_id
+    assert node_display.infer_node_class().__id__ == expected_id

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -28,12 +28,19 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
+        raw_code = raise_if_descriptor(node.code)
+        code_value = None
+
+        if raw_code:
+            code_value = raw_code
+        else:
+            node_file_path = inspect.getfile(node)
+            file_code = read_file_from_path(
+                node_filepath=node_file_path,
+                script_filepath=(raise_if_descriptor(node.filepath)),  # type: ignore
+            )
+            code_value = file_code
 
-        node_file_path = inspect.getfile(node)
-        code = read_file_from_path(
-            node_filepath=node_file_path,
-            script_filepath=(raise_if_descriptor(node.filepath)),  # type: ignore
-        )
         code_inputs = raise_if_descriptor(node.code_inputs)
 
         inputs = [
@@ -50,7 +57,7 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
         code_node_input = create_node_input(
             node_id=node_id,
             input_name="code",
-            value=code,
+            value=code_value,
             display_context=display_context,
             input_id=self.code_input_id,
         )

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/conditional_node.py

@@ -25,10 +25,10 @@ from vellum.workflows.expressions.not_in import NotInExpression
 from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable import ConditionalNode
 from vellum.workflows.types.core import ConditionType, JsonObject
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import NodeInput
 
 _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/final_output_node.py

@@ -3,11 +3,11 @@ from typing import Any, ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 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
 from vellum_ee.workflows.display.nodes.utils import to_kebab_case
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 
 _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -5,11 +5,11 @@ 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
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/merge_node.py

@@ -1,11 +1,11 @@
 from uuid import UUID
-from typing import Any, ClassVar, Generic, List, Optional, Type, TypeVar
+from typing import Any, ClassVar, Generic, List, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable import MergeNode
 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
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import EdgeVellumDisplay
 
 _MergeNodeType = TypeVar("_MergeNodeType", bound=MergeNode)
@@ -14,8 +14,8 @@ _MergeNodeType = TypeVar("_MergeNodeType", bound=MergeNode)
 class BaseMergeNodeDisplay(BaseNodeVellumDisplay[_MergeNodeType], Generic[_MergeNodeType]):
     target_handle_ids: ClassVar[List[UUID]]
 
-    def __init__(self, node: Type[_MergeNodeType]):
-        super().__init__(node)
+    def __init__(self):
+        super().__init__()
         self._target_handle_iterator = 0
 
     def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/search_node.py

@@ -10,11 +10,11 @@ from vellum import (
 from vellum.workflows.nodes.displayable.search_node import SearchNode
 from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import InputVariablePointer, NodeInput
 
 _SearchNodeType = TypeVar("_SearchNodeType", bound=SearchNode)

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -37,6 +37,10 @@ class MyNodeA(BaseNode):
         output: str
 
 
+class MyNodeADisplay(BaseNodeVellumDisplay[MyNodeA]):
+    pass
+
+
 class MyNodeB(BaseNode):
     example = MyNodeA.Outputs.output
     fallback_example = MyNodeA.Outputs.output.coalesce(Inputs.example_workflow_input).coalesce("fallback")
@@ -99,7 +103,7 @@ def test_create_node_input_value_pointer_rules(
                 ),
             },
             node_displays={
-                MyNodeA: BaseNodeVellumDisplay(MyNodeA),
+                MyNodeA: MyNodeADisplay(),
             },
         ),
     )

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -1,13 +1,15 @@
 from uuid import UUID
 from typing import Any, ClassVar, Generic, Optional, TypeVar
 
+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.utils.uuids import uuid4_from_hash
 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
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 
 _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
 
@@ -20,15 +22,19 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
 
         try:
             inner_node = get_wrapped_node(node)
-        except TypeError:
-            raise NotImplementedError(
-                "Unable to serialize Try Nodes that wrap subworkflows containing more than one Node."
-            )
+        except AttributeError:
+            subworkflow = raise_if_descriptor(node.subworkflow)
+            if not isinstance(subworkflow.graph, type) or not issubclass(subworkflow.graph, BaseNode):
+                raise NotImplementedError(
+                    "Unable to serialize Try Nodes that wrap subworkflows containing more than one Node."
+                )
+
+            inner_node = subworkflow.graph
 
         # We need the node display class of the underlying node because
         # it contains the logic for serializing the node and potential display overrides
         node_display_class = get_node_display_class(BaseNodeVellumDisplay, inner_node)
-        node_display = node_display_class(inner_node)
+        node_display = node_display_class()
 
         serialized_node = node_display.serialize(
             display_context,

{vellum_ai-0.11.10 → vellum_ai-0.12.0}/ee/vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -5,8 +5,8 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
 from vellum.workflows.references import NodeReference, OutputReference
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule, primitive_to_vellum_value
 from vellum_ee.workflows.display.vellum import (
     ConstantValuePointer,