vellum-ai 0.14.19__py3-none-any.whl → 0.14.20__py3-none-any.whl

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -12,15 +12,14 @@ from vellum.workflows.edges import Edge
 from vellum.workflows.events.workflow import NodeEventDisplayContext, WorkflowEventDisplayContext
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.utils import get_wrapped_node
+from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.base import (
-    EdgeDisplayOverridesType,
-    EdgeDisplayType,
+    EdgeDisplay,
     EntrypointDisplayOverridesType,
     EntrypointDisplayType,
     StateValueDisplayOverridesType,
@@ -31,11 +30,13 @@ from vellum_ee.workflows.display.base import (
     WorkflowMetaDisplayType,
     WorkflowOutputDisplay,
 )
+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.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
-from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.vellum import EdgeVellumDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 logger = logging.getLogger(__name__)
@@ -50,11 +51,8 @@ class BaseWorkflowDisplay(
         WorkflowInputsDisplayOverridesType,
         StateValueDisplayType,
         StateValueDisplayOverridesType,
-        NodeDisplayType,
         EntrypointDisplayType,
         EntrypointDisplayOverridesType,
-        EdgeDisplayType,
-        EdgeDisplayOverridesType,
     ]
 ):
     # Used to specify the display data for a workflow.
@@ -73,7 +71,7 @@ class BaseWorkflowDisplay(
     output_displays: Dict[BaseDescriptor, WorkflowOutputDisplay] = {}
 
     # Used to explicitly specify display data for a workflow's edges.
-    edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplayOverridesType] = {}
+    edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplay] = {}
 
     # Used to explicitly specify display data for a workflow's ports.
     port_displays: Dict[Port, PortDisplayOverrides] = {}
@@ -94,9 +92,7 @@ class BaseWorkflowDisplay(
                 WorkflowMetaDisplayType,
                 WorkflowInputsDisplayType,
                 StateValueDisplayType,
-                NodeDisplayType,
                 EntrypointDisplayType,
-                EdgeDisplayType,
             ]
         ] = None,
         dry_run: bool = False,
@@ -124,7 +120,7 @@ class BaseWorkflowDisplay(
 
     @property
     @abstractmethod
-    def node_display_base_class(self) -> Type[NodeDisplayType]:
+    def node_display_base_class(self) -> Type[BaseNodeDisplay]:
         pass
 
     def add_error(self, error: Exception) -> None:
@@ -141,7 +137,7 @@ class BaseWorkflowDisplay(
     def _enrich_global_node_output_displays(
         self,
         node: Type[BaseNode],
-        node_display: NodeDisplayType,
+        node_display: BaseNodeDisplay,
         node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]],
     ):
         """This method recursively adds nodes wrapped in decorators to the node_output_displays dictionary."""
@@ -162,7 +158,7 @@ class BaseWorkflowDisplay(
     def _enrich_node_port_displays(
         self,
         node: Type[BaseNode],
-        node_display: NodeDisplayType,
+        node_display: BaseNodeDisplay,
         port_displays: Dict[Port, PortDisplay],
     ):
         """This method recursively adds nodes wrapped in decorators to the port_displays dictionary."""
@@ -178,7 +174,7 @@ class BaseWorkflowDisplay(
 
             port_displays[port] = node_display.get_node_port_display(port)
 
-    def _get_node_display(self, node: Type[BaseNode]) -> NodeDisplayType:
+    def _get_node_display(self, node: Type[BaseNode]) -> BaseNodeDisplay:
         node_display_class = get_node_display_class(self.node_display_base_class, node)
         node_display = node_display_class()
 
@@ -194,9 +190,7 @@ class BaseWorkflowDisplay(
         WorkflowMetaDisplayType,
         WorkflowInputsDisplayType,
         StateValueDisplayType,
-        NodeDisplayType,
         EntrypointDisplayType,
-        EdgeDisplayType,
     ]:
         workflow_display = self._generate_workflow_meta_display()
 
@@ -204,9 +198,9 @@ class BaseWorkflowDisplay(
             copy(self._parent_display_context.global_node_output_displays) if self._parent_display_context else {}
         )
 
-        node_displays: Dict[Type[BaseNode], NodeDisplayType] = {}
+        node_displays: Dict[Type[BaseNode], BaseNodeDisplay] = {}
 
-        global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = (
+        global_node_displays: Dict[Type[BaseNode], BaseNodeDisplay] = (
             copy(self._parent_display_context.global_node_displays) if self._parent_display_context else {}
         )
 
@@ -273,7 +267,7 @@ class BaseWorkflowDisplay(
                 entrypoint, workflow_display, node_displays, overrides=entrypoint_display_overrides
             )
 
-        edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplayType] = {}
+        edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeVellumDisplay] = {}
         for edge in self._workflow.get_edges():
             if edge in edge_displays:
                 continue
@@ -347,7 +341,7 @@ class BaseWorkflowDisplay(
         self,
         entrypoint: Type[BaseNode],
         workflow_display: WorkflowMetaDisplayType,
-        node_displays: Dict[Type[BaseNode], NodeDisplayType],
+        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
         overrides: Optional[EntrypointDisplayOverridesType] = None,
     ) -> EntrypointDisplayType:
         pass
@@ -357,16 +351,6 @@ class BaseWorkflowDisplay(
 
         return WorkflowOutputDisplay(id=output_id, name=output.name)
 
-    @abstractmethod
-    def _generate_edge_display(
-        self,
-        edge: Edge,
-        node_displays: Dict[Type[BaseNode], NodeDisplayType],
-        port_displays: Dict[Port, PortDisplay],
-        overrides: Optional[EdgeDisplayOverridesType] = None,
-    ) -> EdgeDisplayType:
-        pass
-
     def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)
 
@@ -443,9 +427,9 @@ class BaseWorkflowDisplay(
         )
         return display_meta
 
-    def _extract_node_displays(self, node: Type[BaseNode]) -> Dict[Type[BaseNode], NodeDisplayType]:
+    def _extract_node_displays(self, node: Type[BaseNode]) -> Dict[Type[BaseNode], BaseNodeDisplay]:
         node_display = self._get_node_display(node)
-        additional_node_displays: Dict[Type[BaseNode], NodeDisplayType] = {
+        additional_node_displays: Dict[Type[BaseNode], BaseNodeDisplay] = {
             node: node_display,
         }
 
@@ -459,3 +443,47 @@ class BaseWorkflowDisplay(
                     additional_node_displays[node] = display
 
         return additional_node_displays
+
+    def _generate_edge_display(
+        self,
+        edge: Edge,
+        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
+        port_displays: Dict[Port, PortDisplay],
+        overrides: Optional[EdgeDisplay] = None,
+    ) -> EdgeVellumDisplay:
+        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
+        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
+        target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
+
+        return self._generate_edge_display_from_source(
+            source_node_id, source_handle_id, target_node_id, target_handle_id, overrides
+        )
+
+    def _generate_edge_display_from_source(
+        self,
+        source_node_id: UUID,
+        source_handle_id: UUID,
+        target_node_id: UUID,
+        target_handle_id: UUID,
+        overrides: Optional[EdgeDisplay] = None,
+    ) -> EdgeVellumDisplay:
+        edge_id: UUID
+        if overrides:
+            edge_id = overrides.id
+        else:
+            edge_id = uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}")
+
+        return EdgeVellumDisplay(
+            id=edge_id,
+            source_node_id=source_node_id,
+            target_node_id=target_node_id,
+            source_handle_id=source_handle_id,
+            target_handle_id=target_handle_id,
+        )

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -1,14 +1,12 @@
 import logging
 from uuid import UUID
-from typing import Dict, List, Optional, Type, cast
+from typing import Dict, Optional, Type, cast
 
 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.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port
-from vellum.workflows.ports import Port
 from vellum.workflows.references import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
@@ -16,12 +14,9 @@ from vellum.workflows.types.generics import WorkflowType
 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.types import PortDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import (
-    EdgeVellumDisplay,
-    EdgeVellumDisplayOverrides,
     EntrypointVellumDisplay,
     EntrypointVellumDisplayOverrides,
     NodeDisplayData,
@@ -46,11 +41,8 @@ class VellumWorkflowDisplay(
         WorkflowInputsVellumDisplayOverrides,
         StateValueVellumDisplay,
         StateValueVellumDisplayOverrides,
-        BaseNodeDisplay,
         EntrypointVellumDisplay,
         EntrypointVellumDisplayOverrides,
-        EdgeVellumDisplay,
-        EdgeVellumDisplayOverrides,
     ]
 ):
     node_display_base_class = BaseNodeDisplay
@@ -99,14 +91,16 @@ class VellumWorkflowDisplay(
         edges: JsonArray = []
 
         # Add a single synthetic node for the workflow entrypoint
+        entrypoint_node_id = self.display_context.workflow_display.entrypoint_node_id
+        entrypoint_node_source_handle_id = self.display_context.workflow_display.entrypoint_node_source_handle_id
         nodes.append(
             {
-                "id": str(self.display_context.workflow_display.entrypoint_node_id),
+                "id": str(entrypoint_node_id),
                 "type": "ENTRYPOINT",
                 "inputs": [],
                 "data": {
                     "label": "Entrypoint Node",
-                    "source_handle_id": str(self.display_context.workflow_display.entrypoint_node_source_handle_id),
+                    "source_handle_id": str(entrypoint_node_source_handle_id),
                 },
                 "display_data": self.display_context.workflow_display.entrypoint_node_display.dict(),
                 "base": None,
@@ -245,25 +239,37 @@ class VellumWorkflowDisplay(
             )
 
         # Add an edge for each edge in the workflow
-        all_edge_displays: List[EdgeVellumDisplay] = [
-            # Create a synthetic edge from the synthetic entrypoint node to each actual entrypoint
-            *[
-                entrypoint_display.edge_display
-                for entrypoint_display in self.display_context.entrypoint_displays.values()
-            ],
-            # Include the concrete edges in the workflow
-            *self.display_context.edge_displays.values(),
-        ]
+        for target_node, entrypoint_display in self.display_context.entrypoint_displays.items():
+            unadorned_target_node = get_unadorned_node(target_node)
+            target_node_display = self.display_context.node_displays[unadorned_target_node]
+            edges.append(
+                {
+                    "id": str(entrypoint_display.edge_display.id),
+                    "source_node_id": str(entrypoint_node_id),
+                    "source_handle_id": str(entrypoint_node_source_handle_id),
+                    "target_node_id": str(target_node_display.node_id),
+                    "target_handle_id": str(target_node_display.get_trigger_id()),
+                    "type": "DEFAULT",
+                }
+            )
+
+        for (source_node_port, target_node), edge_display in self.display_context.edge_displays.items():
+            unadorned_source_node_port = get_unadorned_port(source_node_port)
+            unadorned_target_node = get_unadorned_node(target_node)
+
+            source_node_port_display = self.display_context.port_displays[unadorned_source_node_port]
+            target_node_display = self.display_context.node_displays[unadorned_target_node]
 
-        for edge_display in all_edge_displays:
             edges.append(
                 {
                     "id": str(edge_display.id),
-                    "source_node_id": str(edge_display.source_node_id),
-                    "source_handle_id": str(edge_display.source_handle_id),
-                    "target_node_id": str(edge_display.target_node_id),
-                    "target_handle_id": str(edge_display.target_handle_id),
-                    "type": edge_display.type,
+                    "source_node_id": str(source_node_port_display.node_id),
+                    "source_handle_id": str(source_node_port_display.id),
+                    "target_node_id": str(target_node_display.node_id),
+                    "target_handle_id": str(
+                        target_node_display.get_target_handle_id_by_source_node_id(source_node_port_display.node_id)
+                    ),
+                    "type": "DEFAULT",
                 }
             )
 
@@ -357,61 +363,10 @@ class VellumWorkflowDisplay(
         entrypoint_target = get_unadorned_node(entrypoint)
         target_node_display = node_displays[entrypoint_target]
         target_node_id = target_node_display.node_id
-        if isinstance(target_node_display, BaseNodeVellumDisplay):
-            target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(entrypoint_node_id)
-        else:
-            target_handle_id = target_node_display.get_trigger_id()
+        target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(entrypoint_node_id)
 
         edge_display = self._generate_edge_display_from_source(
             entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides
         )
 
         return EntrypointVellumDisplay(id=entrypoint_id, edge_display=edge_display)
-
-    def _generate_edge_display(
-        self,
-        edge: Edge,
-        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
-        port_displays: Dict[Port, PortDisplay],
-        overrides: Optional[EdgeVellumDisplayOverrides] = None,
-    ) -> EdgeVellumDisplay:
-        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
-        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
-
-        if isinstance(target_node_display, BaseNodeVellumDisplay):
-            target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
-        else:
-            target_handle_id = target_node_display.get_trigger_id()
-
-        return self._generate_edge_display_from_source(
-            source_node_id, source_handle_id, target_node_id, target_handle_id, overrides
-        )
-
-    def _generate_edge_display_from_source(
-        self,
-        source_node_id: UUID,
-        source_handle_id: UUID,
-        target_node_id: UUID,
-        target_handle_id: UUID,
-        overrides: Optional[EdgeVellumDisplayOverrides] = None,
-    ) -> EdgeVellumDisplay:
-        edge_id: UUID
-        if overrides:
-            edge_id = overrides.id
-        else:
-            edge_id = uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}")
-
-        return EdgeVellumDisplay(
-            id=edge_id,
-            source_node_id=source_node_id,
-            target_node_id=target_node_id,
-            source_handle_id=source_handle_id,
-            target_handle_id=target_handle_id,
-        )

vellum_ee/workflows/server/virtual_file_loader.py

@@ -1,42 +1,72 @@
 import importlib
+import re
+from typing import Optional
 
 
 class VirtualFileLoader(importlib.abc.Loader):
-    def __init__(self, code: str, is_package: bool):
-        self.code = code
-        self.is_package = is_package
+    def __init__(self, files: dict[str, str], namespace: str):
+        self.files = files
+        self.namespace = namespace
 
     def create_module(self, spec):
         return None  # use default module creation
 
     def exec_module(self, module):
-        if not self.is_package or self.code:
-            exec(self.code, module.__dict__)
+        module_info = self._resolve_module(module.__spec__.origin)
 
+        if module_info:
+            file_path, code = module_info
+            compiled = compile(code, file_path, "exec")
+            exec(compiled, module.__dict__)
 
-class VirtualFileFinder(importlib.abc.MetaPathFinder, importlib.abc.Loader):
-    def __init__(self, files: dict[str, str], namespace: str):
-        self.files = files
-        self.namespace = namespace
+    def get_source(self, fullname):
+        """
+        `inspect` module uses this method to get the source code of a module.
+        """
+
+        module_info = self._resolve_module(fullname)
+        if module_info:
+            return module_info[1]
+
+        return None
+
+    def _resolve_module(self, fullname: str) -> Optional[tuple[str, str]]:
+        file_path = self._get_file_path(fullname)
+        code = self._get_code(file_path)
+
+        if code is not None:
+            return file_path, code
 
-    def find_spec(self, fullname, path, target=None):
-        # Do the namespacing on the fly to avoid having to copy the file dict
-        prefixed_name = fullname if fullname.startswith(self.namespace) else f"{self.namespace}.{fullname}"
+        if not file_path.endswith("__init__.py"):
+            file_path = re.sub(r"\.py$", "/__init__.py", file_path)
+            code = self._get_code(file_path)
 
-        key_name = "__init__" if fullname == self.namespace else fullname.replace(f"{self.namespace}.", "")
+            if code is not None:
+                return file_path, code
+
+        return None
 
-        files_key = f"{key_name.replace('.', '/')}.py"
-        if self.files.get(files_key) is None:
-            files_key = f"{key_name.replace('.', '/')}/__init__.py"
+    def _get_file_path(self, fullname):
+        return f"{fullname.replace('.', '/')}.py"
 
-        file = self.files.get(files_key)
-        is_package = "__init__" in files_key
+    def _get_code(self, file_path):
+        file_key_name = re.sub(r"^" + re.escape(self.namespace) + r"/", "", file_path)
+        return self.files.get(file_key_name)
+
+
+class VirtualFileFinder(importlib.abc.MetaPathFinder, importlib.abc.Loader):
+    def __init__(self, files: dict[str, str], namespace: str):
+        self.loader = VirtualFileLoader(files, namespace)
 
-        if file is not None:
+    def find_spec(self, fullname: str, path, target=None):
+        module_info = self.loader._resolve_module(fullname)
+        if module_info:
+            file_path, _ = module_info
+            is_package = file_path.endswith("__init__.py")
             return importlib.machinery.ModuleSpec(
-                prefixed_name,
-                VirtualFileLoader(file, is_package),
-                origin=prefixed_name,
+                fullname,
+                self.loader,
+                origin=fullname,
                 is_package=is_package,
             )
         return None

vellum_ee/workflows/tests/test_server.py

@@ -1,4 +1,11 @@
+import sys
+from uuid import uuid4
+from typing import Type, cast
+
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes import BaseNode
+from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
 
 def test_load_workflow_event_display_context():
@@ -8,3 +15,57 @@ def test_load_workflow_event_display_context():
     # We are actually just ensuring there are no circular dependencies when
     # our Workflow Server imports this class.
     assert issubclass(WorkflowEventDisplayContext, UniversalBaseModel)
+
+
+def test_load_from_module__lazy_reference_in_file_loader():
+    # GIVEN a workflow module with a node containing a lazy reference
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.start_node import StartNode
+
+class Workflow(BaseWorkflow):
+    graph = StartNode
+""",
+        "nodes/__init__.py": """\
+from .start_node import StartNode
+
+__all__ = [
+    "StartNode",
+]
+""",
+        "nodes/start_node.py": """\
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.references import LazyReference
+
+class StartNode(BaseNode):
+    foo = LazyReference(lambda: StartNode.Outputs.bar)
+
+    class Outputs(BaseNode.Outputs):
+        bar = str
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    sys.meta_path.append(VirtualFileFinder(files, namespace))
+
+    # WHEN the workflow is loaded
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    workflow = Workflow()
+
+    # THEN the workflow is successfully initialized
+    assert workflow
+
+    # AND the graph is just a BaseNode
+    # ideally this would be true, but the loader uses a different BaseNode class definition than
+    # the one in this test module.
+    # assert isinstance(workflow.graph, BaseNode)
+    start_node = cast(Type[BaseNode], workflow.graph)
+    assert start_node.__bases__ == (BaseNode,)
+
+    # AND the lazy reference has the correct name
+    assert start_node.foo.instance
+    assert start_node.foo.instance.name == "StartNode.Outputs.bar"