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

vellum/workflows/nodes/displayable/bases/search_node.py

@@ -0,0 +1,90 @@
+from decimal import Decimal
+from uuid import UUID
+from typing import ClassVar, Generic, List, Optional, Union
+
+from vellum import (
+    NotFoundError,
+    SearchFiltersRequest,
+    SearchRequestOptionsRequest,
+    SearchResponse,
+    SearchResult,
+    SearchResultMergingRequest,
+    SearchWeightsRequest,
+)
+from vellum.core import ApiError, RequestOptions
+
+from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.types.generics import StateType
+
+DEFAULT_SEARCH_WEIGHTS = 0.8
+DEFAULT_SEARCH_LIMIT = 8
+
+
+def get_default_results_merging() -> SearchResultMergingRequest:
+    return SearchResultMergingRequest(enabled=True)
+
+
+class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
+    """
+    Used to perform a hybrid search against a Document Index in Vellum.
+
+    document_index: Union[UUID, str] - Either the Document Index's UUID or its name.
+    query: str - The query to search for.
+    options: Optional[SearchRequestOptionsRequest] = None - The request options to use for the search
+    request_options: Optional[RequestOptions] = None - The request options to use for the search
+    """
+
+    # The query to search for.
+    query: ClassVar[str]
+
+    # The Document Index to Search against. Identified by either its UUID or its name.
+    document_index: ClassVar[Union[UUID, str]]
+
+    # Ideally we could reuse node descriptors to derive other node descriptor values. Two action items are
+    # blocking us from doing so in this use case:
+    # 1. Node Descriptor resolution during runtime - https://app.shortcut.com/vellum/story/4781
+    # 2. Math operations between descriptors - https://app.shortcut.com/vellum/story/4782
+    # search_weights = DEFAULT_SEARCH_WEIGHTS
+    options = SearchRequestOptionsRequest(
+        limit=DEFAULT_SEARCH_LIMIT,
+        weights=SearchWeightsRequest(
+            semantic_similarity=DEFAULT_SEARCH_WEIGHTS,
+            keywords=float(Decimal("1.0") - Decimal(str(DEFAULT_SEARCH_WEIGHTS))),
+        ),
+        result_merging=get_default_results_merging(),
+        filters=SearchFiltersRequest(
+            external_ids=None,
+            metadata=None,
+        ),
+    )
+
+    request_options: Optional[RequestOptions] = None
+
+    class Outputs(BaseOutputs):
+        results: List[SearchResult]
+
+    def _perform_search(self) -> SearchResponse:
+        try:
+            return self._context.vellum_client.search(
+                query=self.query,
+                index_id=str(self.document_index) if isinstance(self.document_index, UUID) else None,
+                index_name=self.document_index if isinstance(self.document_index, str) else None,
+                options=self.options,
+            )
+        except NotFoundError:
+            raise NodeException(
+                message=f"Document Index '{self.document_index}' not found",
+                code=VellumErrorCode.INVALID_INPUTS,
+            )
+        except ApiError:
+            raise NodeException(
+                message=f"An error occurred while searching against Document Index '{self.document_index}'",  # noqa: E501
+                code=VellumErrorCode.INTERNAL_ERROR,
+            )
+
+    def run(self) -> Outputs:
+        response = self._perform_search()
+        return self.Outputs(results=response.results)

vellum/workflows/nodes/displayable/code_execution_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import CodeExecutionNode
+
+__all__ = [
+    "CodeExecutionNode",
+]

vellum/workflows/nodes/displayable/code_execution_node/node.py

@@ -0,0 +1,197 @@
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Sequence, Tuple, Type, TypeVar, cast, get_args
+
+from vellum import (
+    ArrayInput,
+    ChatHistoryInput,
+    ChatMessage,
+    CodeExecutionPackage,
+    CodeExecutionRuntime,
+    CodeExecutorInput,
+    ErrorInput,
+    FunctionCall,
+    FunctionCallInput,
+    JsonInput,
+    NumberInput,
+    SearchResult,
+    SearchResultsInput,
+    StringInput,
+    VellumError,
+    VellumValue,
+)
+from vellum.core import RequestOptions
+
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.bases.base import BaseNodeMeta
+from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.types.core import EntityInputsInterface, VellumSecret
+from vellum.workflows.types.generics import StateType
+from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+
+_OutputType = TypeVar("_OutputType")
+
+
+# TODO: Consolidate all dynamic output metaclasses
+# https://app.shortcut.com/vellum/story/5533
+class _CodeExecutionNodeMeta(BaseNodeMeta):
+    def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        parent = super().__new__(mcs, name, bases, dct)
+
+        # We use the compiled class to infer the output type for the Outputs.result descriptor.
+        if not isinstance(parent, _CodeExecutionNodeMeta):
+            raise ValueError("CodeExecutionNode must be created with the CodeExecutionNodeMeta metaclass")
+
+        parent.__dict__["Outputs"].__annotations__["result"] = parent.get_output_type()
+        return parent
+
+    def get_output_type(cls) -> Type:
+        original_base = get_original_base(cls)
+        all_args = get_args(original_base)
+
+        if len(all_args) < 2 or isinstance(all_args[1], TypeVar):
+            return str
+        else:
+            return all_args[1]
+
+
+class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], metaclass=_CodeExecutionNodeMeta):
+    """
+    Used to execute an arbitrary script. This node exists to be backwards compatible with
+    Vellum's Code Execution Node, and for most cases, you should extend from `BaseNode` directly.
+
+    filepath: str - The path to the script to execute.
+    code_inputs: EntityInputsInterface - The inputs for the custom script.
+    output_type: VellumVariableType = "STRING" - The type of the output from the custom script.
+    runtime: CodeExecutionRuntime = "PYTHON_3_12" - The runtime to use for the custom script.
+    packages: Optional[Sequence[CodeExecutionPackageRequest]] = None - The packages to use for the custom script.
+    request_options: Optional[RequestOptions] = None - The request options to use for the custom script.
+    """
+
+    filepath: ClassVar[str]
+
+    code_inputs: ClassVar[EntityInputsInterface]
+    runtime: CodeExecutionRuntime = "PYTHON_3_11_6"
+    packages: Optional[Sequence[CodeExecutionPackage]] = None
+
+    request_options: Optional[RequestOptions] = None
+
+    class Outputs(BaseOutputs):
+        # We use our mypy plugin to override the _OutputType with the actual output type
+        # for downstream references to this output.
+        result: _OutputType  # type: ignore[valid-type]
+        log: str
+
+    def run(self) -> Outputs:
+        input_values = self._compile_code_inputs()
+        expected_output_type = primitive_type_to_vellum_variable_type(self.__class__.get_output_type())
+        code_execution = self._context.vellum_client.execute_code(
+            input_values=input_values,
+            code=self._resolve_code(),
+            runtime=self.runtime,
+            output_type=expected_output_type,
+            packages=self.packages or [],
+            request_options=self.request_options,
+        )
+
+        if code_execution.output.type != expected_output_type:
+            raise NodeException(
+                code=VellumErrorCode.INVALID_OUTPUTS,
+                message=f"Expected an output of type '{expected_output_type}', received '{code_execution.output.type}'",
+            )
+
+        return self.Outputs(result=code_execution.output.value, log=code_execution.log)
+
+    def _compile_code_inputs(self) -> List[CodeExecutorInput]:
+        # TODO: We may want to consolidate with prompt deployment input compilation
+        # https://app.shortcut.com/vellum/story/4117
+
+        compiled_inputs: List[CodeExecutorInput] = []
+
+        for input_name, input_value in self.code_inputs.items():
+            if isinstance(input_value, str):
+                compiled_inputs.append(
+                    StringInput(
+                        name=input_name,
+                        value=str(input_value),
+                    )
+                )
+            elif isinstance(input_value, VellumSecret):
+                compiled_inputs.append(
+                    # TODO: Expose a VellumSecret type from the Vellum SDK
+                    # https://app.shortcut.com/vellum/story/4785
+                    {  # type: ignore[arg-type]
+                        "name": input_name,
+                        "type": "SECRET",
+                        "value": input_value.name,
+                    }
+                )
+            elif isinstance(input_value, list):
+                if all(isinstance(message, ChatMessage) for message in input_value):
+                    compiled_inputs.append(
+                        ChatHistoryInput(
+                            name=input_name,
+                            value=cast(List[ChatMessage], input_value),
+                        )
+                    )
+                elif all(isinstance(message, SearchResult) for message in input_value):
+                    compiled_inputs.append(
+                        SearchResultsInput(
+                            name=input_name,
+                            value=cast(List[SearchResult], input_value),
+                        )
+                    )
+                else:
+                    compiled_inputs.append(
+                        ArrayInput(
+                            name=input_name,
+                            value=cast(List[VellumValue], input_value),
+                        )
+                    )
+            elif isinstance(input_value, dict):
+                compiled_inputs.append(
+                    JsonInput(
+                        name=input_name,
+                        value=cast(Dict[str, Any], input_value),
+                    )
+                )
+            elif isinstance(input_value, float):
+                compiled_inputs.append(
+                    NumberInput(
+                        name=input_name,
+                        value=input_value,
+                    )
+                )
+            elif isinstance(input_value, FunctionCall):
+                compiled_inputs.append(
+                    FunctionCallInput(
+                        name=input_name,
+                        value=cast(FunctionCall, input_value),
+                    )
+                )
+            elif isinstance(input_value, VellumError):
+                compiled_inputs.append(
+                    ErrorInput(
+                        name=input_name,
+                        value=cast(VellumError, input_value),
+                    )
+                )
+            else:
+                raise NodeException(
+                    message=f"Unrecognized input type for input '{input_name}'",
+                    code=VellumErrorCode.INVALID_INPUTS,
+                )
+
+        return compiled_inputs
+
+    def _resolve_code(self) -> str:
+        code = read_file_from_path(self.filepath)
+        if not code:
+            raise NodeException(
+                message=f"Filepath '{self.filepath}' does not exist",
+                code=VellumErrorCode.INVALID_INPUTS,
+            )
+
+        return code

vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/main.py

@@ -0,0 +1,3 @@
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -0,0 +1,111 @@
+import os
+
+from vellum import CodeExecutorResponse, NumberVellumValue, StringInput
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
+from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.state.base import BaseState, StateMeta
+
+
+def test_run_workflow__happy_path(vellum_client):
+    """Confirm that CodeExecutionNodes output the expected text and results when run."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        filepath = fixture
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": 5, "log": "hello"}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            StringInput(name="word", value="hello"),
+        ],
+        code="""\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[],
+        request_options=None,
+    )
+
+
+def test_run_workflow__vellum_secret(vellum_client):
+    """Confirm that CodeExecutionNodes can use Vellum Secrets"""
+
+    # GIVEN a node that subclasses CodeExecutionNode that references a Vellum Secret
+    class State(BaseState):
+        pass
+
+    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        filepath = fixture
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "token": VellumSecretReference("OPENAI_API_KEY"),
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="",
+        output=NumberVellumValue(value=0),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": 0, "log": ""}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            {"name": "token", "type": "SECRET", "value": "OPENAI_API_KEY"},
+        ],
+        code="""\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[],
+        request_options=None,
+    )

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -0,0 +1,10 @@
+import os
+from typing import Union
+
+
+def read_file_from_path(filepath: str) -> Union[str, None]:
+    if not os.path.exists(filepath):
+        return None
+
+    with open(filepath) as file:
+        return file.read()

vellum/workflows/nodes/displayable/conditional_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import ConditionalNode
+
+__all__ = [
+    "ConditionalNode",
+]

vellum/workflows/nodes/displayable/conditional_node/node.py

@@ -0,0 +1,25 @@
+from typing import Iterable
+
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.ports.node_ports import NodePorts
+from vellum.workflows.ports.port import Port
+from vellum.workflows.ports.utils import validate_ports
+from vellum.workflows.state.base import BaseState
+
+
+class ConditionalNode(BaseNode):
+    """
+    Used to conditionally determine which port to invoke next. This node exists to be backwards compatible with
+    Vellum's Conditional Node, and for most cases, you should extend `BaseNode.Ports` directly.
+    """
+
+    class Ports(NodePorts):
+        def __call__(self, outputs: BaseOutputs, state: BaseState) -> Iterable[Port]:
+            all_ports = [port for port in self.__class__]
+            enforce_single_invoked_port = validate_ports(all_ports)
+
+            if not enforce_single_invoked_port:
+                raise ValueError("Conditional nodes must have exactly one if port")
+
+            return super().__call__(outputs, state)

vellum/workflows/nodes/displayable/final_output_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import FinalOutputNode
+
+__all__ = [
+    "FinalOutputNode",
+]

vellum/workflows/nodes/displayable/final_output_node/node.py

@@ -0,0 +1,43 @@
+from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args
+
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.bases.base import BaseNodeMeta
+from vellum.workflows.types.generics import StateType
+from vellum.workflows.types.utils import get_original_base
+
+_OutputType = TypeVar("_OutputType")
+
+
+# TODO: Consolidate all dynamic output metaclasses
+# https://app.shortcut.com/vellum/story/5533
+class _FinalOutputNodeMeta(BaseNodeMeta):
+    def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        parent = super().__new__(mcs, name, bases, dct)
+
+        # We use the compiled class to infer the output type for the Outputs.value descriptor.
+        if not isinstance(parent, _FinalOutputNodeMeta):
+            raise ValueError("CodeExecutionNode must be created with the CodeExecutionNodeMeta metaclass")
+
+        parent.__dict__["Outputs"].__annotations__["value"] = parent.get_output_type()
+        return parent
+
+    def get_output_type(cls) -> Type:
+        original_base = get_original_base(cls)
+        all_args = get_args(original_base)
+
+        if len(all_args) < 2 or isinstance(all_args[1], TypeVar):
+            return str
+        else:
+            return all_args[1]
+
+
+class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], metaclass=_FinalOutputNodeMeta):
+    """
+    Used to directly reference the output of another node.
+    This provides backward compatibility with Vellum's Final Output Node.
+    """
+
+    class Outputs(BaseNode.Outputs):
+        # We use our mypy plugin to override the _OutputType with the actual output type
+        # for downstream references to this output.
+        value: _OutputType  # type: ignore[valid-type]

vellum/workflows/nodes/displayable/guardrail_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import GuardrailNode
+
+__all__ = [
+    "GuardrailNode",
+]

vellum/workflows/nodes/displayable/guardrail_node/node.py

@@ -0,0 +1,97 @@
+from uuid import UUID
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Union, cast
+
+from vellum import ChatHistoryInput, ChatMessage, JsonInput, MetricDefinitionInput, NumberInput, StringInput
+from vellum.core import RequestOptions
+
+from vellum.workflows.constants import LATEST_RELEASE_TAG
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.types.core import EntityInputsInterface
+from vellum.workflows.types.generics import StateType
+
+
+class GuardrailNode(BaseNode[StateType], Generic[StateType]):
+    """
+    Used to execute a Metric Definition and surface a float output representing the score.
+
+    metric_definition: Union[UUID, str] - Either the Metric Definition's UUID or its name.
+    metric_inputs: EntityInputsInterface - The inputs for the Metric
+    release_tag: str - The release tag to use for the Metric
+    request_options: Optional[RequestOptions] - The request options to use for the Metric
+    """
+
+    metric_definition: ClassVar[Union[UUID, str]]
+
+    metric_inputs: ClassVar[EntityInputsInterface]
+    release_tag: str = LATEST_RELEASE_TAG
+
+    request_options: Optional[RequestOptions] = None
+
+    class Outputs(BaseOutputs):
+        score: float
+
+    def run(self) -> Outputs:
+        metric_execution = self._context.vellum_client.metric_definitions.execute_metric_definition(
+            self.metric_definition if isinstance(self.metric_definition, str) else str(self.metric_definition),
+            inputs=self._compile_metric_inputs(),
+            release_tag=self.release_tag,
+            request_options=self.request_options,
+        )
+
+        metric_outputs = {output.name: output.value for output in metric_execution.outputs}
+
+        score = metric_outputs.get("score")
+        if not isinstance(score, float):
+            raise NodeException(
+                message="Metric execution must have one output named 'score' with type 'float'",
+                code=VellumErrorCode.INVALID_OUTPUTS,
+            )
+
+        metric_outputs.pop("score")
+        return self.Outputs(score=score, **metric_outputs)
+
+    def _compile_metric_inputs(self) -> List[MetricDefinitionInput]:
+        # TODO: We may want to consolidate with prompt deployment input compilation
+        # https://app.shortcut.com/vellum/story/4117
+
+        compiled_inputs: List[MetricDefinitionInput] = []
+
+        for input_name, input_value in self.metric_inputs.items():
+            if isinstance(input_value, str):
+                compiled_inputs.append(
+                    StringInput(
+                        name=input_name,
+                        value=input_value,
+                    )
+                )
+            elif isinstance(input_value, list) and all(isinstance(message, ChatMessage) for message in input_value):
+                compiled_inputs.append(
+                    ChatHistoryInput(
+                        name=input_name,
+                        value=cast(List[ChatMessage], input_value),
+                    )
+                )
+            elif isinstance(input_value, dict):
+                compiled_inputs.append(
+                    JsonInput(
+                        name=input_name,
+                        value=cast(Dict[str, Any], input_value),
+                    )
+                )
+            elif isinstance(input_value, float):
+                compiled_inputs.append(
+                    NumberInput(
+                        name=input_name,
+                        value=input_value,
+                    )
+                )
+            else:
+                raise NodeException(
+                    message=f"Unrecognized input type for input '{input_name}'",
+                    code=VellumErrorCode.INVALID_INPUTS,
+                )
+
+        return compiled_inputs

vellum/workflows/nodes/displayable/inline_prompt_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import InlinePromptNode
+
+__all__ = [
+    "InlinePromptNode",
+]

vellum/workflows/nodes/displayable/inline_prompt_node/node.py

@@ -0,0 +1,41 @@
+from typing import Iterator
+
+from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.displayable.bases import BaseInlinePromptNode as BaseInlinePromptNode
+from vellum.workflows.outputs import BaseOutput
+from vellum.workflows.types.generics import StateType
+
+
+class InlinePromptNode(BaseInlinePromptNode[StateType]):
+    """
+    Used to execute an Inline Prompt and surface a string output for convenience.
+
+    prompt_inputs: EntityInputsInterface - The inputs for the Prompt
+    ml_model: str - Either the ML Model's UUID or its name.
+    blocks: List[PromptBlockRequest] - The blocks that make up the Prompt
+    parameters: PromptParameters - The parameters for the Prompt
+    expand_meta: Optional[AdHocExpandMetaRequest] - Set of expandable execution fields to include in the response
+    """
+
+    class Outputs(BaseInlinePromptNode.Outputs):
+        text: str
+
+    def run(self) -> Iterator[BaseOutput]:
+        outputs = yield from self._process_prompt_event_stream()
+        if not outputs:
+            raise NodeException(
+                message="Expected to receive outputs from Prompt",
+                code=VellumErrorCode.INTERNAL_ERROR,
+            )
+
+        string_output = next((output for output in outputs if output.type == "STRING"), None)
+        if not string_output or string_output.value is None:
+            output_types = {output.type for output in outputs}
+            is_plural = len(output_types) > 1
+            raise NodeException(
+                message=f"Expected to receive a non-null string output from Prompt. Only found outputs of type{'s' if is_plural else ''}: {', '.join(output_types)}",  # noqa: E501
+                code=VellumErrorCode.INTERNAL_ERROR,
+            )
+
+        yield BaseOutput(name="text", value=string_output.value)

vellum/workflows/nodes/displayable/merge_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import MergeNode
+
+__all__ = [
+    "MergeNode",
+]

vellum/workflows/nodes/displayable/merge_node/node.py

@@ -0,0 +1,10 @@
+from vellum.workflows.nodes.bases import BaseNode
+
+
+class MergeNode(BaseNode):
+    """
+    Used to merge the control flow of multiple nodes into a single node. This node exists to be backwards compatible
+    with Vellum's Merge Node, and for most cases, you should extend from `BaseNode.Trigger` directly.
+    """
+
+    pass