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

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/local_workflow/display/workflow.py

@@ -1,7 +1,7 @@
 from uuid import UUID
 
+from vellum_ee.workflows.display.base import EdgeDisplay
 from vellum_ee.workflows.display.vellum import (
-    EdgeVellumDisplayOverrides,
     EntrypointVellumDisplayOverrides,
     NodeDisplayData,
     NodeDisplayPosition,
@@ -36,13 +36,11 @@ class WorkflowDisplay(VellumWorkflowDisplay[Workflow]):
     entrypoint_displays = {
         TemplatingNode: EntrypointVellumDisplayOverrides(
             id=UUID("0bf86989-13f2-438c-ab9c-d172e5771d31"),
-            edge_display=EdgeVellumDisplayOverrides(id=UUID("38532a0e-9432-4ed2-8a34-48a29fd6984d")),
+            edge_display=EdgeDisplay(id=UUID("38532a0e-9432-4ed2-8a34-48a29fd6984d")),
         )
     }
     edge_displays = {
-        (TemplatingNode.Ports.default, FinalOutput): EdgeVellumDisplayOverrides(
-            id=UUID("417c56a4-cdc1-4f9d-a10c-b535163f51e8")
-        )
+        (TemplatingNode.Ports.default, FinalOutput): EdgeDisplay(id=UUID("417c56a4-cdc1-4f9d-a10c-b535163f51e8"))
     }
     output_displays = {
         Workflow.Outputs.final_output: WorkflowOutputVellumDisplayOverrides(

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"