vellum-ai 0.9.16rc2__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

vellum_cli/tests/test_config.py

@@ -0,0 +1,100 @@
+import os
+import tempfile
+from uuid import uuid4
+
+from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_config
+
+
+def test_load_vellum_cli_config__pyproject_toml_only():
+    # GIVEN a pyproject.toml file with vellum config
+    cwd = tempfile.mkdtemp()
+    with open(os.path.join(cwd, "pyproject.toml"), "w") as f:
+        f.write(
+            """[[tool.vellum.workflows]]
+module = "module1.workflow1"
+"""
+        )
+
+    # WHEN the config is loaded
+    config = load_vellum_cli_config(cwd)
+
+    # THEN the config is loaded correctly
+    assert config == VellumCliConfig(
+        workflows=[WorkflowConfig(module="module1.workflow1")],
+        version="1.0",
+    )
+
+
+def test_load_vellum_cli_config__pyproject_toml_and_lockfile():
+    # GIVEN a pyproject.toml file with vellum config
+    cwd = tempfile.mkdtemp()
+    with open(os.path.join(cwd, "pyproject.toml"), "w") as f:
+        f.write(
+            """[[tool.vellum.workflows]]
+module = "module1.workflow1"
+"""
+        )
+
+    # AND a lockfile
+    workflow_sandbox_id = uuid4()
+    with open(os.path.join(cwd, "vellum.lock.json"), "w") as f:
+        f.write(
+            f"""{{
+  "version": "1.0",
+  "workflows": [
+    {{
+      "workflow_sandbox_id": "{workflow_sandbox_id}",
+      "module": "module1.workflow1"
+    }}
+  ]
+}}
+"""
+        )
+
+    # WHEN the config is loaded
+    config = load_vellum_cli_config(cwd)
+
+    # THEN the config is loaded correctly
+    assert config == VellumCliConfig(
+        workflows=[WorkflowConfig(module="module1.workflow1", workflow_sandbox_id=str(workflow_sandbox_id))],
+        version="1.0",
+    )
+
+
+def test_load_vellum_cli_config__pyproject_toml_and_lockfile__different_modules():
+    # GIVEN a pyproject.toml file with vellum config
+    cwd = tempfile.mkdtemp()
+    with open(os.path.join(cwd, "pyproject.toml"), "w") as f:
+        f.write(
+            """[[tool.vellum.workflows]]
+module = "module1.workflow1"
+"""
+        )
+
+    # AND a lockfile
+    workflow_sandbox_id = uuid4()
+    with open(os.path.join(cwd, "vellum.lock.json"), "w") as f:
+        f.write(
+            f"""{{
+  "version": "1.0",
+  "workflows": [
+    {{
+      "workflow_sandbox_id": "{workflow_sandbox_id}",
+      "module": "module2.workflow2"
+    }}
+  ]
+}}
+"""
+        )
+
+    # WHEN the config is loaded
+    config = load_vellum_cli_config(cwd)
+
+    # THEN the config is loaded correctly
+    assert config == VellumCliConfig(
+        workflows=[
+            WorkflowConfig(module="module1.workflow1"),
+            WorkflowConfig(module="module2.workflow2", workflow_sandbox_id=str(workflow_sandbox_id)),
+        ],
+        version="1.0",
+    )

vellum_cli/tests/test_pull.py

@@ -0,0 +1,152 @@
+import pytest
+import io
+import os
+import shutil
+import tempfile
+from uuid import uuid4
+import zipfile
+from typing import Generator, Tuple
+
+import tomli_w
+
+from vellum_cli.pull import pull_command
+
+
+def zip_file_map(file_map: dict[str, str]) -> bytes:
+    # Create an in-memory bytes buffer to store the zip
+    zip_buffer = io.BytesIO()
+
+    # Create zip file and add files from file_map
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+        for filename, content in file_map.items():
+            zip_file.writestr(filename, content)
+
+    # Get the bytes from the buffer
+    zip_bytes = zip_buffer.getvalue()
+    zip_buffer.close()
+
+    return zip_bytes
+
+
+@pytest.fixture
+def mock_module() -> Generator[Tuple[str, str], None, None]:
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+    module = "examples.mock"
+
+    with open(os.path.join(temp_dir, "pyproject.toml"), "wb") as f:
+        tomli_w.dump(
+            {
+                "tool": {
+                    "vellum": {
+                        "workflows": [
+                            {
+                                "module": module,
+                                "workflow_sandbox_id": str(uuid4()),
+                            }
+                        ]
+                    }
+                }
+            },
+            f,
+        )
+
+    yield temp_dir, module
+
+    os.chdir(current_dir)
+    shutil.rmtree(temp_dir)
+
+
+def test_pull(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    temp_dir, module = mock_module
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command
+    pull_command(module)
+
+    # THEN the workflow.py file is written to the module directory
+    assert os.path.exists(os.path.join(temp_dir, *module.split("."), "workflow.py"))
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py")) as f:
+        assert f.read() == "print('hello')"
+
+
+def test_pull__remove_missing_files(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    temp_dir, module = mock_module
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
+
+    # AND there is already a different file in the module directory
+    other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
+    os.makedirs(os.path.dirname(other_file_path), exist_ok=True)
+    with open(other_file_path, "w") as f:
+        f.write("print('hello')")
+
+    # WHEN the user runs the pull command
+    pull_command(module)
+
+    # THEN the workflow.py file is written to the module directory
+    assert os.path.exists(os.path.join(temp_dir, *module.split("."), "workflow.py"))
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py")) as f:
+        assert f.read() == "print('hello')"
+
+    # AND the other_file.py file is deleted
+    assert not os.path.exists(other_file_path)
+
+
+def test_pull__remove_missing_files__ignore_pattern(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    temp_dir, module = mock_module
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
+
+    # AND there is already a different file in the module directory
+    other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
+    os.makedirs(os.path.dirname(other_file_path), exist_ok=True)
+    with open(other_file_path, "w") as f:
+        f.write("print('hello')")
+
+    # AND there is already a test file
+    test_file_path = os.path.join(temp_dir, *module.split("."), "tests", "test_workflow.py")
+    os.makedirs(os.path.dirname(test_file_path), exist_ok=True)
+    with open(test_file_path, "w") as f:
+        f.write("print('hello')")
+
+    # AND the ignore pattern is set to tests
+    with open(os.path.join(temp_dir, "pyproject.toml"), "wb") as f:
+        tomli_w.dump(
+            {
+                "tool": {
+                    "vellum": {
+                        "workflows": [
+                            {
+                                "module": module,
+                                "workflow_sandbox_id": str(uuid4()),
+                                "ignore": "tests/*",
+                            }
+                        ]
+                    }
+                }
+            },
+            f,
+        )
+
+    # WHEN the user runs the pull command
+    pull_command(module)
+
+    # THEN the workflow.py file is written to the module directory
+    assert os.path.exists(os.path.join(temp_dir, *module.split("."), "workflow.py"))
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py")) as f:
+        assert f.read() == "print('hello')"
+
+    # AND the other_file.py file is deleted
+    assert not os.path.exists(other_file_path)
+
+    # AND the tests/test_workflow.py file is untouched
+    assert os.path.exists(test_file_path)

vellum_ee/workflows/display/base.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+from uuid import UUID
+from typing import TypeVar
+
+
+@dataclass
+class WorkflowMetaDisplayOverrides:
+    pass
+
+
+@dataclass
+class WorkflowMetaDisplay(WorkflowMetaDisplayOverrides):
+    pass
+
+
+WorkflowMetaDisplayType = TypeVar("WorkflowMetaDisplayType", bound=WorkflowMetaDisplay)
+WorkflowMetaDisplayOverridesType = TypeVar("WorkflowMetaDisplayOverridesType", bound=WorkflowMetaDisplayOverrides)
+
+
+@dataclass
+class WorkflowInputsDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class WorkflowInputsDisplay(WorkflowInputsDisplayOverrides):
+    pass
+
+
+WorkflowInputsDisplayType = TypeVar("WorkflowInputsDisplayType", bound=WorkflowInputsDisplay)
+WorkflowInputsDisplayOverridesType = TypeVar("WorkflowInputsDisplayOverridesType", bound=WorkflowInputsDisplayOverrides)
+
+
+@dataclass
+class EdgeDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class EdgeDisplay(EdgeDisplayOverrides):
+    pass
+
+
+EdgeDisplayType = TypeVar("EdgeDisplayType", bound=EdgeDisplay)
+EdgeDisplayOverridesType = TypeVar("EdgeDisplayOverridesType", bound=EdgeDisplayOverrides)
+
+
+@dataclass
+class EntrypointDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class EntrypointDisplay(EntrypointDisplayOverrides):
+    pass
+
+
+EntrypointDisplayType = TypeVar("EntrypointDisplayType", bound=EntrypointDisplay)
+EntrypointDisplayOverridesType = TypeVar("EntrypointDisplayOverridesType", bound=EntrypointDisplayOverrides)
+
+
+@dataclass
+class WorkflowOutputDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class WorkflowOutputDisplay(WorkflowOutputDisplayOverrides):
+    pass
+
+
+WorkflowOutputDisplayType = TypeVar("WorkflowOutputDisplayType", bound=WorkflowOutputDisplay)
+WorkflowOutputDisplayOverridesType = TypeVar("WorkflowOutputDisplayOverridesType", bound=WorkflowOutputDisplayOverrides)

vellum_ee/workflows/display/nodes/__init__.py

@@ -0,0 +1,4 @@
+# flake8: noqa: F401
+
+# Force an import to ensure that all display classes are registered with the BaseNodeDisplay registry
+from .vellum import *  # noqa: F401

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -0,0 +1,116 @@
+from functools import cached_property
+import inspect
+from uuid import UUID
+from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, get_args
+
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.vellum import CodeResourceDefinition, NodeDefinition
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.ports import Port
+from vellum.workflows.references import OutputReference
+from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.generics import NodeType
+from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.names import pascal_to_title_case
+
+if TYPE_CHECKING:
+    from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+_NodeDisplayAttrType = TypeVar("_NodeDisplayAttrType")
+
+
+class BaseNodeDisplay(Generic[NodeType]):
+    output_display: Dict[OutputReference, NodeOutputDisplay] = {}
+    port_displays: Dict[Port, PortDisplayOverrides] = {}
+
+    # Used to store the mapping between node types and their display classes
+    _node_display_registry: Dict[Type[NodeType], Type["BaseNodeDisplay"]] = {}
+
+    def __init__(self, node: Type[NodeType]):
+        self._node = node
+
+    def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
+        raise NotImplementedError(f"Serialization for nodes of type {self._node.__name__} is not supported.")
+
+    def get_definition(self) -> NodeDefinition:
+        node = self._node
+        node_definition = NodeDefinition(
+            name=node.__name__,
+            module=node.__module__.split("."),
+            bases=[
+                CodeResourceDefinition(
+                    name=base.__name__,
+                    module=base.__module__.split("."),
+                )
+                for base in node.__bases__
+            ],
+        )
+        return node_definition
+
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
+        explicit_display = self.output_display.get(output)
+        if explicit_display:
+            return explicit_display
+
+        return NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name)
+
+    def get_node_port_display(self, port: Port) -> PortDisplay:
+        overrides = self.port_displays.get(port)
+
+        port_id: UUID
+        if overrides:
+            port_id = overrides.id
+        else:
+            port_id = uuid4_from_hash(f"{self.node_id}|ports|{port.name}")
+
+        return PortDisplay(id=port_id, node_id=self.node_id)
+
+    @classmethod
+    def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Type["BaseNodeDisplay"]:
+        return cls._node_display_registry[node_class]
+
+    @cached_property
+    def node_id(self) -> UUID:
+        """Can be overridden as a class attribute to specify a custom node id."""
+        return uuid4_from_hash(self._node.__qualname__)
+
+    @cached_property
+    def label(self) -> str:
+        """Can be overridden as a class attribute to specify a custom label."""
+        return pascal_to_title_case(self._node.__name__)
+
+    @classmethod
+    def _get_explicit_node_display_attr(
+        cls,
+        attribute: str,
+        attribute_type: Type[_NodeDisplayAttrType],
+    ) -> Optional[_NodeDisplayAttrType]:
+        node_display_attribute: Optional[_NodeDisplayAttrType] = getattr(cls, attribute, None)
+
+        if node_display_attribute is None:
+            return None
+
+        if isinstance(node_display_attribute, attribute_type):
+            return node_display_attribute
+
+        raise ValueError(f"Node {cls.__name__} must define an explicit {attribute} of type {attribute_type.__name__}.")
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        original_base = get_original_base(cls)
+        node_class = get_args(original_base)[0]
+        if isinstance(node_class, TypeVar):
+            bounded_class = node_class.__bound__
+            if inspect.isclass(bounded_class) and issubclass(bounded_class, BaseNode):
+                cls._node_display_registry[bounded_class] = cls
+        elif issubclass(node_class, BaseNode):
+            if has_wrapped_node(node_class):
+                wrapped_node = get_wrapped_node(node_class)
+                if wrapped_node._is_wrapped_node:
+                    cls._node_display_registry[wrapped_node] = cls
+                    return
+
+            cls._node_display_registry[node_class] = cls

vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -0,0 +1,36 @@
+from uuid import UUID
+from typing import ClassVar, Dict, Optional
+
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.types import PortDisplay
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.vellum import NodeDisplayData
+from vellum.workflows.ports import Port
+from vellum.workflows.types.generics import NodeType
+
+
+class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
+    # Used to explicitly set display data for a node
+    display_data: ClassVar[Optional[NodeDisplayData]] = None
+
+    # Used to explicitly set the target handle id for a node
+    target_handle_id: ClassVar[Optional[UUID]] = None
+
+    # Used to explicitly set the node input ids by name for a node
+    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+
+    def _get_node_display_uuid(self, attribute: str) -> UUID:
+        explicit_value = self._get_explicit_node_display_attr(attribute, UUID)
+        return explicit_value if explicit_value else uuid4_from_hash(f"{self.node_id}|{attribute}")
+
+    def get_display_data(self) -> NodeDisplayData:
+        explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
+        return explicit_value if explicit_value else NodeDisplayData()
+
+    def get_target_handle_id(self) -> UUID:
+        return self._get_node_display_uuid("target_handle_id")
+
+    def get_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
+        default_port = self._node.Ports.default
+        default_port_display = port_displays[default_port]
+        return default_port_display.id

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -0,0 +1,25 @@
+from typing import Optional, Type
+
+from vellum_ee.workflows.display.types import NodeDisplayType
+from vellum.workflows.types.generics import NodeType
+
+
+def get_node_display_class(
+    base_class: Type[NodeDisplayType], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
+) -> Type[NodeDisplayType]:
+    try:
+        node_display_class = base_class.get_from_node_display_registry(node_class)
+    except KeyError:
+        try:
+            return get_node_display_class(
+                base_class, node_class.__bases__[0], node_class if root_node_class is None else root_node_class
+            )
+        except IndexError:
+            return base_class
+
+    if not issubclass(node_display_class, base_class):
+        raise TypeError(
+            f"Expected to find a subclass of '{base_class.__name__}' for node class '{node_class.__name__}'"
+        )
+
+    return node_display_class

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -0,0 +1,47 @@
+import pytest
+from uuid import UUID
+
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum.workflows.nodes.bases import BaseNode
+
+
+@pytest.fixture
+def node_with_implicit_properties():
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        pass
+
+    expected_id = UUID("ace7f746-4fe6-45c7-8207-fc8a4d0c7f6f")
+
+    return MyNode, MyNodeDisplay, expected_id
+
+
+@pytest.fixture
+def node_with_explicit_properties():
+    explicit_id = UUID("a422f67a-1d37-43f0-bdfc-1e4618c9496d")
+
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        node_id = explicit_id
+
+    return MyNode, MyNodeDisplay, explicit_id
+
+
+@pytest.fixture(
+    params=[
+        "node_with_implicit_properties",
+        "node_with_explicit_properties",
+    ]
+)
+def node_info(request):
+    return request.getfixturevalue(request.param)
+
+
+def test_get_id(node_info):
+    node, node_display, expected_id = node_info
+
+    assert node_display(node).node_id == expected_id

vellum_ee/workflows/display/nodes/types.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from uuid import UUID
+
+
+@dataclass
+class NodeOutputDisplay:
+    id: UUID
+    name: str
+
+
+@dataclass
+class PortDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class PortDisplay(PortDisplayOverrides):
+    node_id: UUID

vellum_ee/workflows/display/nodes/utils.py

@@ -0,0 +1,33 @@
+import re
+from typing import Optional, TypeVar, Union, overload
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.references.node import NodeReference
+
+_T = TypeVar("_T")
+
+
+@overload
+def raise_if_descriptor(node_attr: BaseDescriptor[_T]) -> _T: ...
+
+
+@overload
+def raise_if_descriptor(node_attr: _T) -> _T: ...
+
+
+def raise_if_descriptor(node_attr: Union[NodeReference[_T], _T]) -> Optional[_T]:
+    if not isinstance(node_attr, NodeReference):
+        raise AttributeError(f"Expected to find a Node descriptor, but found '{node_attr.__class__.__name__}'")
+
+    return node_attr.instance
+
+
+def to_kebab_case(string: str) -> str:
+    """Converts a string from Capital Case to kebab-case."""
+    string = re.sub(r"(\s|_|-)+", " ", string)
+    string = re.sub(
+        r"[A-Z]{2,}(?=[A-Z][a-z]+[0-9]*|\b)|[A-Z]?[a-z]+[0-9]*|[A-Z]|[0-9]+",
+        lambda mo: " " + mo.group(0).lower(),
+        string,
+    )
+    return "-".join(string.split())

vellum_ee/workflows/display/nodes/vellum/__init__.py

@@ -0,0 +1,32 @@
+from .api_node import BaseAPINodeDisplay
+from .code_execution_node import BaseCodeExecutionNodeDisplay
+from .conditional_node import BaseConditionalNodeDisplay
+from .final_output_node import BaseFinalOutputNodeDisplay
+from .guardrail_node import BaseGuardrailNodeDisplay
+from .inline_prompt_node import BaseInlinePromptNodeDisplay
+from .inline_subworkflow_node import BaseInlineSubworkflowNodeDisplay
+from .map_node import BaseMapNodeDisplay
+from .merge_node import BaseMergeNodeDisplay
+from .prompt_deployment_node import BasePromptDeploymentNodeDisplay
+from .search_node import BaseSearchNodeDisplay
+from .subworkflow_deployment_node import BaseSubworkflowDeploymentNodeDisplay
+from .templating_node import BaseTemplatingNodeDisplay
+from .try_node import BaseTryNodeDisplay
+
+# All node display classes must be imported here to be registered in BaseNodeDisplay's node display registry
+__all__ = [
+    "BaseCodeExecutionNodeDisplay",
+    "BaseConditionalNodeDisplay",
+    "BaseGuardrailNodeDisplay",
+    "BaseInlinePromptNodeDisplay",
+    "BaseInlineSubworkflowNodeDisplay",
+    "BaseAPINodeDisplay",
+    "BaseMapNodeDisplay",
+    "BaseMergeNodeDisplay",
+    "BaseSearchNodeDisplay",
+    "BaseSubworkflowDeploymentNodeDisplay",
+    "BaseTemplatingNodeDisplay",
+    "BasePromptDeploymentNodeDisplay",
+    "BaseFinalOutputNodeDisplay",
+    "BaseTryNodeDisplay",
+]