vellum-ai 1.8.4__py3-none-any.whl → 1.8.6__py3-none-any.whl

vellum_ee/workflows/tests/test_server.py

@@ -34,13 +34,6 @@ 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
@@ -98,11 +91,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = CodeExecutionNode.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/code_execution_node/__init__.py": """
 from typing import Any
@@ -174,11 +162,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = CodeExecutionNode.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -250,11 +233,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = CodeExecutionNode.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -326,11 +304,6 @@ class Workflow(BaseWorkflow):
 
     class Outputs(BaseWorkflow.Outputs):
        final_output = Subworkflow.Outputs.result
-""",
-        "nodes/__init__.py": """
-from .subworkflow import Subworkflow
-
-__all__ = ["Subworkflow"]
 """,
         "nodes/subworkflow/__init__.py": """
 from vellum.workflows.nodes.displayable import InlineSubworkflowNode
@@ -339,11 +312,6 @@ from .workflow import SubworkflowWorkflow
 
 class Subworkflow(InlineSubworkflowNode):
     subworkflow = SubworkflowWorkflow
-""",
-        "nodes/subworkflow/nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/subworkflow/nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -424,11 +392,6 @@ class Workflow(BaseWorkflow):
     
     class Outputs(BaseWorkflow.Outputs):  # noqa: W293
        results = MapNode.Outputs.final_output
-""",
-        "nodes/__init__.py": """
-from .map_node import MapNode
-
-__all__ = ["MapNode"]
 """,
         "nodes/map_node/__init__.py": """
 from vellum.workflows.nodes.core.map_node import MapNode as BaseMapNode
@@ -439,11 +402,6 @@ class MapNode(BaseMapNode):
     items = ["foo", "bar", "baz"]
     subworkflow = MapNodeWorkflow
     max_concurrency = 4
-""",
-        "nodes/map_node/nodes/__init__.py": """
-from .code_execution_node import CodeExecutionNode
-
-__all__ = ["CodeExecutionNode"]
 """,
         "nodes/map_node/nodes/code_execution_node/__init__.py": """
 from typing import Union
@@ -520,7 +478,6 @@ from .nodes.broken_node import BrokenNode
 class Workflow(BaseWorkflow):
     graph = BrokenNode
 """,
-        "nodes/__init__.py": "",
         "nodes/broken_node.py": """\
 from vellum.workflows.nodes import BaseNode
 
@@ -559,7 +516,6 @@ from .nodes.broken_node import BrokenNode
 class Workflow(BaseWorkflow):
     graph = BrokenNode
 """,
-        "nodes/__init__.py": "",
         "nodes/broken_node.py": """\
 from vellum.workflows.nodes import BaseNode
 
@@ -819,19 +775,9 @@ class TestSubworkflowDeploymentNode(SubworkflowDeploymentNode):
     files = {
         "__init__.py": "",
         "workflow.py": parent_workflow_code,
-        "nodes/__init__.py": """
-from .subworkflow_deployment_node import TestSubworkflowDeploymentNode
-
-__all__ = ["TestSubworkflowDeploymentNode"]
-""",
         "nodes/subworkflow_deployment_node.py": parent_node_code,
         f"{expected_prefix}/__init__.py": "",
         f"{expected_prefix}/workflow.py": mock_workflow_code,
-        f"{expected_prefix}/nodes/__init__.py": """
-from .test_node import TestNode
-
-__all__ = ["TestNode"]
-""",
         f"{expected_prefix}/nodes/test_node.py": test_node_code,
     }
 
@@ -919,11 +865,6 @@ class TestSubworkflowDeploymentNode(SubworkflowDeploymentNode):
     subworkflow_files = {
         "__init__.py": "",
         "workflow.py": mock_workflow_code,
-        "nodes/__init__.py": """
-from .test_node import TestNode
-
-__all__ = ["TestNode"]
-""",
         "nodes/test_node.py": test_node_code,
     }
 
@@ -931,11 +872,6 @@ __all__ = ["TestNode"]
         "__init__.py": "",
         "inputs.py": inputs_code,
         "workflow.py": parent_workflow_code,
-        "nodes/__init__.py": """
-from .subworkflow_deployment_node import TestSubworkflowDeploymentNode
-
-__all__ = ["TestSubworkflowDeploymentNode"]
-""",
         "nodes/subworkflow_deployment_node.py": parent_node_code,
     }
 

vellum_ee/workflows/tests/test_virtual_files.py

@@ -80,3 +80,54 @@ def test_base_class_dynamic_import(files):
         assert instance.id
     except Exception as e:
         pytest.fail(f"Failed to create an instance of BaseClass: {e}")
+
+
+def test_display_directory_not_auto_generated():
+    """
+    Test that the top-level display directory is NOT auto-generated with empty __init__.py.
+    Display directories typically have specific __init__.py content (e.g., "from .workflow import *")
+    that should not be replaced with empty auto-generated files.
+    """
+    # GIVEN a workflow with display/workflow.py but NO display/__init__.py
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.bases import BaseNode
+
+class StartNode(BaseNode):
+    pass
+
+class Workflow(BaseWorkflow):
+    graph = StartNode
+""",
+        "display/workflow.py": """\
+from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+
+class WorkflowDisplay(BaseWorkflowDisplay):
+    pass
+""",
+        # Note: NO "display/__init__.py" in files dict
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    sys.meta_path.append(VirtualFileFinder(files, namespace))
+
+    try:
+        # WHEN we try to resolve display/__init__.py
+        import importlib.util
+
+        spec = importlib.util.find_spec(f"{namespace}.display")
+
+        # THEN the spec should be None because we don't want to auto-generate display/__init__.py
+        # If the spec exists, it means an empty __init__.py was auto-generated (BAD)
+        assert spec is None, (
+            "display directory should NOT have auto-generated __init__.py. "
+            "Display directories require specific __init__.py content that shouldn't be empty."
+        )
+
+    finally:
+        # Clean up
+        sys.meta_path = [finder for finder in sys.meta_path if not isinstance(finder, VirtualFileFinder)]