vellum-ai 1.7.10__py3-none-any.whl → 1.7.12__py3-none-any.whl

vellum_ee/workflows/server/virtual_file_loader.py

@@ -4,11 +4,14 @@ import re
 import sys
 from typing import Optional
 
+from vellum.workflows.loaders.base import BaseWorkflowFinder
+
 
 class VirtualFileLoader(importlib.abc.Loader):
-    def __init__(self, files: dict[str, str], namespace: str):
+    def __init__(self, files: dict[str, str], namespace: str, source_module: Optional[str] = None):
         self.files = files
         self.namespace = namespace
+        self.source_module = source_module
 
     def create_module(self, spec: ModuleSpec):
         """
@@ -65,9 +68,17 @@ class VirtualFileLoader(importlib.abc.Loader):
         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)
+class VirtualFileFinder(BaseWorkflowFinder):
+    def __init__(self, files: dict[str, str], namespace: str, source_module: Optional[str] = None):
+        self.loader = VirtualFileLoader(files, namespace, source_module)
+        self.source_module = source_module
+        self.namespace = namespace
+
+    def format_error_message(self, error_message: str) -> str:
+        """Format error message by replacing namespace with source_module."""
+        if self.source_module and self.namespace in error_message:
+            return error_message.replace(self.namespace, self.source_module)
+        return error_message
 
     def find_spec(self, fullname: str, path, target=None):
         module_info = self.loader._resolve_module(fullname)

vellum_ee/workflows/tests/test_serialize_module.py

@@ -96,6 +96,54 @@ def test_serialize_module_includes_additional_files():
     assert "CONSTANT_VALUE" in additional_files["utils/constants.py"]
 
 
+def test_serialize_module_with_pydantic_array():
+    """
+    Test that serialize_module correctly serializes arrays of Pydantic models in workflow inputs.
+
+    This test verifies that when a workflow has inputs containing a List[PydanticModel],
+    the serialization properly converts the Pydantic models to JSON format.
+    """
+    module_path = "tests.workflows.pydantic_array_serialization"
+
+    # WHEN we serialize it
+    result = BaseWorkflowDisplay.serialize_module(module_path)
+
+    assert hasattr(result, "exec_config")
+    assert hasattr(result, "errors")
+    assert isinstance(result.exec_config, dict)
+    assert isinstance(result.errors, list)
+
+    input_variables = result.exec_config["input_variables"]
+    assert len(input_variables) == 1
+
+    items_input = input_variables[0]
+    assert items_input["key"] == "items"
+    assert items_input["type"] == "JSON"
+    # TODO: In the future, this should be a custom type based on an OpenAPI schema (important-comment)
+
+    assert result.dataset is not None
+    assert isinstance(result.dataset, list)
+    assert len(result.dataset) == 2
+
+    first_scenario = result.dataset[0]
+    assert first_scenario["label"] == "Scenario 1"
+    assert "items" in first_scenario["inputs"]
+    items = first_scenario["inputs"]["items"]
+    assert isinstance(items, list)
+    assert len(items) == 3
+    assert items[0]["name"] == "item1"
+    assert items[0]["value"] == 10
+    assert items[0]["is_active"] is True
+
+    second_scenario = result.dataset[1]
+    assert second_scenario["label"] == "Custom Test"
+    assert "items" in second_scenario["inputs"]
+    test_items = second_scenario["inputs"]["items"]
+    assert len(test_items) == 2
+    assert test_items[0]["name"] == "test1"
+    assert test_items[0]["value"] == 100
+
+
 def test_serialize_module__with_invalid_nested_set_graph(temp_module_path):
     """
     Tests that serialize_module raises a clear user-facing exception for workflows with nested sets in graph attribute.

vellum_ee/workflows/tests/test_server.py

@@ -584,6 +584,111 @@ class BrokenNode(BaseNode):
     assert "UndefinedClass" in error_message or "not defined" in error_message
 
 
+def test_load_from_module__module_not_found_error():
+    """
+    Tests that a ModuleNotFoundError raises WorkflowInitializationException with user-facing message.
+    """
+    # GIVEN a workflow module that imports a non-existent module
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .non_existent_module import SomeClass
+
+class Workflow(BaseWorkflow):
+    graph = None
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    sys.meta_path.append(VirtualFileFinder(files, namespace, source_module="test"))
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly and show source_module instead of namespace
+    error_message = str(exc_info.value)
+    assert error_message == "Workflow module not found: No module named 'test.non_existent_module'"
+
+
+def test_load_from_module__module_not_found_error_with_external_package():
+    """
+    Tests that when ModuleNotFoundError occurs for an external package (not containing the namespace),
+    the exception includes vellum_on_error_action set to CREATE_CUSTOM_IMAGE in raw_data.
+    """
+
+    # GIVEN a workflow module that imports a non-existent external package
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+import some_external_package
+
+class Workflow(BaseWorkflow):
+    graph = None
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace, source_module="test")
+    sys.meta_path.append(finder)
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly
+    error_message = str(exc_info.value)
+    assert "Workflow module not found:" in error_message
+    assert "some_external_package" in error_message
+
+    assert exc_info.value.raw_data is not None
+    assert exc_info.value.raw_data["vellum_on_error_action"] == "CREATE_CUSTOM_IMAGE"
+
+
+def test_load_from_module__module_not_found_error_with_internal_package():
+    """
+    Tests that when ModuleNotFoundError occurs for an internal module (containing the namespace),
+    the exception does NOT include vellum_on_error_action in raw_data.
+    """
+
+    # GIVEN a workflow module that imports a non-existent internal module
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .non_existent_module import SomeClass
+
+class Workflow(BaseWorkflow):
+    graph = None
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace, source_module="test")
+    sys.meta_path.append(finder)
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly
+    error_message = str(exc_info.value)
+    assert "Workflow module not found:" in error_message
+
+    assert exc_info.value.raw_data is None
+
+
 def test_serialize_module__tool_calling_node_with_single_tool():
     """Test that serialize_module works with a tool calling node that has a single tool."""