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

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"