vellum-ai 1.4.1__py3-none-any.whl → 1.5.0__py3-none-any.whl

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

@@ -5,6 +5,7 @@ from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
 from vellum.workflows.nodes.utils import cast_to_output_type
 from vellum.workflows.ports import NodePorts
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.types.utils import get_original_base
@@ -27,6 +28,7 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
             **annotations,
             "value": parent.get_output_type(),
         }
+
         return parent
 
     def get_output_type(cls) -> Type:
@@ -38,6 +40,63 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
         else:
             return all_args[1]
 
+    def __validate__(cls) -> None:
+        cls._validate_output_type_consistency(cls)
+
+    @classmethod
+    def _validate_output_type_consistency(mcs, cls: Type) -> None:
+        """
+        Validates that the declared output type of FinalOutputNode matches
+        the type of the descriptor assigned to the 'value' attribute in its Outputs class.
+
+        Raises ValueError if there's a type mismatch.
+        """
+        if not hasattr(cls, "Outputs"):
+            return
+
+        outputs_class = cls.Outputs
+        if not hasattr(outputs_class, "value"):
+            return
+
+        declared_output_type = cls.get_output_type()
+        value_descriptor = None
+
+        if "value" in outputs_class.__dict__:
+            value_descriptor = outputs_class.__dict__["value"]
+        else:
+            value_descriptor = getattr(outputs_class, "value")
+
+        if isinstance(value_descriptor, OutputReference):
+            descriptor_types = value_descriptor.types
+
+            type_mismatch = True
+            for descriptor_type in descriptor_types:
+                if descriptor_type == declared_output_type:
+                    type_mismatch = False
+                    break
+                try:
+                    if issubclass(descriptor_type, declared_output_type) or issubclass(
+                        declared_output_type, descriptor_type
+                    ):
+                        type_mismatch = False
+                        break
+                except TypeError:
+                    # Handle cases where types aren't classes (e.g., Union)
+                    if str(descriptor_type) == str(declared_output_type):
+                        type_mismatch = False
+                        break
+
+            if type_mismatch:
+                declared_type_name = getattr(declared_output_type, "__name__", str(declared_output_type))
+                descriptor_type_names = [getattr(t, "__name__", str(t)) for t in descriptor_types]
+
+                raise ValueError(
+                    f"Output type mismatch in {cls.__name__}: "
+                    f"FinalOutputNode is declared with output type '{declared_type_name}' "
+                    f"but the 'value' descriptor has type(s) {descriptor_type_names}. "
+                    f"The output descriptor type must match the declared FinalOutputNode output type."
+                )
+
 
 class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], metaclass=_FinalOutputNodeMeta):
     """

vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py

@@ -2,10 +2,11 @@ import pytest
 
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.nodes.displayable.inline_prompt_node import InlinePromptNode
 from vellum.workflows.state.base import BaseState
 
 
-def test_final_output_node__mismatched_output_type():
+def test_final_output_node__mismatched_output_type_should_raise_exception_when_ran():
     # GIVEN a FinalOutputNode with a mismatched output type
     class StringOutputNode(FinalOutputNode[BaseState, str]):
         class Outputs(FinalOutputNode.Outputs):
@@ -18,3 +19,41 @@ def test_final_output_node__mismatched_output_type():
 
     # THEN an error is raised
     assert str(exc_info.value) == "Expected an output of type 'str', but received 'dict'"
+
+
+def test_final_output_node__mismatched_output_type_should_raise_exception():
+    # GIVEN a FinalOutputNode declared with list output type but has a string value type
+    class Output(FinalOutputNode[BaseState, list]):
+        """Output the extracted invoice line items as an array of objects."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = InlinePromptNode.Outputs.text
+
+    # WHEN attempting to validate the node class
+    # THEN a ValueError should be raised during validation
+    with pytest.raises(ValueError) as exc_info:
+        Output.__validate__()
+
+    # AND the error message should indicate the type mismatch
+    assert (
+        str(exc_info.value)
+        == "Output type mismatch in Output: FinalOutputNode is declared with output type 'list' but "
+        "the 'value' descriptor has type(s) ['str']. The output descriptor type must match the "
+        "declared FinalOutputNode output type."
+    )
+
+
+def test_final_output_node__matching_output_type_should_pass_validation():
+    # GIVEN a FinalOutputNode declared with correct matching types
+    class CorrectOutput(FinalOutputNode[BaseState, str]):
+        """Output with correct type matching."""
+
+        class Outputs(FinalOutputNode.Outputs):
+            value = InlinePromptNode.Outputs.text
+
+    # WHEN attempting to validate the node class
+    # THEN validation should pass without raising an exception
+    try:
+        CorrectOutput.__validate__()
+    except ValueError:
+        pytest.fail("Validation should not raise an exception for correct type matching")

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

@@ -2,6 +2,7 @@ from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set,
 
 from vellum import ChatMessage, PromptBlock
 from vellum.client.types.prompt_parameters import PromptParameters
+from vellum.client.types.prompt_settings import PromptSettings
 from vellum.prompts.constants import DEFAULT_PROMPT_PARAMETERS
 from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.errors.types import WorkflowErrorCode
@@ -47,6 +48,7 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
     prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
     parameters: PromptParameters = DEFAULT_PROMPT_PARAMETERS
     max_prompt_iterations: ClassVar[Optional[int]] = 5
+    settings: ClassVar[Optional[Union[PromptSettings, Dict[str, Any]]]] = None
 
     class Outputs(BaseOutputs):
         """
@@ -150,6 +152,7 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
             max_prompt_iterations=self.max_prompt_iterations,
             process_parameters_method=process_parameters_method,
             process_blocks_method=process_blocks_method,
+            settings=self.settings,
         )
 
         # Create the router node (handles routing logic only)

vellum/workflows/nodes/displayable/tool_calling_node/tests/test_utils.py

@@ -5,6 +5,7 @@ from vellum.client.types.chat_message_prompt_block import ChatMessagePromptBlock
 from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
 from vellum.client.types.initiated_execute_prompt_event import InitiatedExecutePromptEvent
 from vellum.client.types.plain_text_prompt_block import PlainTextPromptBlock
+from vellum.client.types.prompt_settings import PromptSettings
 from vellum.client.types.rich_text_prompt_block import RichTextPromptBlock
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
@@ -250,3 +251,66 @@ def test_get_mcp_tool_name_snake_case():
 
     result = get_mcp_tool_name(mcp_tool)
     assert result == "github_server__create_repository"
+
+
+def test_create_tool_prompt_node_settings_dict_stream_disabled(vellum_adhoc_prompt_client):
+    # GIVEN settings provided as dict with stream disabled
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[],
+        prompt_inputs=None,
+        parameters=DEFAULT_PROMPT_PARAMETERS,
+        max_prompt_iterations=1,
+        settings={"stream_enabled": False},
+    )
+
+    # AND the API mocks
+    def generate_non_stream_response(*args, **kwargs):
+        return FulfilledExecutePromptEvent(execution_id=str(uuid4()), outputs=[StringVellumValue(value="ok")])
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.side_effect = generate_non_stream_response
+
+    # WHEN we run the node
+    node_instance = tool_prompt_node()
+    list(node_instance.run())
+
+    # THEN the node should have called the API correctly
+    assert node_instance.settings is not None
+    assert node_instance.settings.stream_enabled is False
+    assert vellum_adhoc_prompt_client.adhoc_execute_prompt.call_count == 1
+    assert vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.call_count == 0
+
+
+def test_create_tool_prompt_node_settings_model_stream_enabled(vellum_adhoc_prompt_client):
+    # GIVEN settings provided as PromptSettings with stream enabled
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[],
+        prompt_inputs=None,
+        parameters=DEFAULT_PROMPT_PARAMETERS,
+        max_prompt_iterations=1,
+        settings=PromptSettings(stream_enabled=True),
+    )
+
+    # AND the API mocks
+    def generate_stream_events(*args, **kwargs):
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(execution_id=execution_id, outputs=[StringVellumValue(value="ok")]),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_stream_events
+
+    # WHEN we run the node
+    node_instance = tool_prompt_node()
+    list(node_instance.run())
+
+    # THEN the node should have called the API correctly
+    assert node_instance.settings is not None
+    assert node_instance.settings.stream_enabled is True
+    assert vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.call_count == 1
+    assert vellum_adhoc_prompt_client.adhoc_execute_prompt.call_count == 0

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

@@ -9,9 +9,9 @@ from vellum.client.types.array_chat_message_content import ArrayChatMessageConte
 from vellum.client.types.array_chat_message_content_item import ArrayChatMessageContentItem
 from vellum.client.types.function_call_chat_message_content import FunctionCallChatMessageContent
 from vellum.client.types.function_call_chat_message_content_value import FunctionCallChatMessageContentValue
-from vellum.client.types.function_definition import FunctionDefinition
 from vellum.client.types.prompt_output import PromptOutput
 from vellum.client.types.prompt_parameters import PromptParameters
+from vellum.client.types.prompt_settings import PromptSettings
 from vellum.client.types.string_chat_message_content import StringChatMessageContent
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -31,7 +31,13 @@ from vellum.workflows.ports.port import Port
 from vellum.workflows.state import BaseState
 from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior, Tool, ToolBase
-from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition, MCPServer, MCPToolDefinition
+from vellum.workflows.types.definition import (
+    ComposioToolDefinition,
+    DeploymentDefinition,
+    MCPServer,
+    MCPToolDefinition,
+    VellumIntegrationToolDefinition,
+)
 from vellum.workflows.types.generics import is_workflow_class
 from vellum.workflows.utils.functions import compile_mcp_tool_definition, get_mcp_tool_name
 
@@ -274,36 +280,6 @@ class ElseNode(BaseNode[ToolCallingState]):
         return self.Outputs()
 
 
-def _hydrate_composio_tool_definition(tool_def: ComposioToolDefinition) -> FunctionDefinition:
-    """Hydrate a ComposioToolDefinition with detailed information from the Composio API.
-
-    Args:
-        tool_def: The basic ComposioToolDefinition to enhance
-
-    Returns:
-        FunctionDefinition with detailed parameters and description
-    """
-    try:
-        composio_service = ComposioService()
-        tool_details = composio_service.get_tool_by_slug(tool_def.action)
-
-        # Create a FunctionDefinition directly with proper field extraction
-        return FunctionDefinition(
-            name=tool_def.name,
-            description=tool_details.get("description", tool_def.description),
-            parameters=tool_details.get("input_parameters", {}),
-        )
-
-    except Exception as e:
-        # If hydration fails (including no API key), log and return basic function definition
-        logger.warning(f"Failed to enhance Composio tool '{tool_def.action}': {e}")
-        return FunctionDefinition(
-            name=tool_def.name,
-            description=tool_def.description,
-            parameters={},
-        )
-
-
 def create_tool_prompt_node(
     ml_model: str,
     blocks: List[Union[PromptBlock, Dict[str, Any]]],
@@ -313,17 +289,10 @@ def create_tool_prompt_node(
     max_prompt_iterations: Optional[int] = None,
     process_parameters_method: Optional[Callable] = None,
     process_blocks_method: Optional[Callable] = None,
+    settings: Optional[Union[PromptSettings, Dict[str, Any]]] = None,
 ) -> Type[ToolPromptNode]:
     if functions and len(functions) > 0:
-        prompt_functions: List[Union[Tool, FunctionDefinition]] = []
-
-        for function in functions:
-            if isinstance(function, ComposioToolDefinition):
-                # Get Composio tool details and hydrate the function definition
-                enhanced_function = _hydrate_composio_tool_definition(function)
-                prompt_functions.append(enhanced_function)
-            else:
-                prompt_functions.append(function)
+        prompt_functions: List[Tool] = functions
     else:
         prompt_functions = []
 
@@ -359,6 +328,13 @@ def create_tool_prompt_node(
         ),
     }
 
+    # Normalize settings to PromptSettings if provided as a dict
+    normalized_settings: Optional[PromptSettings]
+    if isinstance(settings, dict):
+        normalized_settings = PromptSettings.model_validate(settings)
+    else:
+        normalized_settings = settings
+
     node = cast(
         Type[ToolPromptNode],
         type(
@@ -371,6 +347,7 @@ def create_tool_prompt_node(
                 "prompt_inputs": node_prompt_inputs,
                 "parameters": parameters,
                 "max_prompt_iterations": max_prompt_iterations,
+                "settings": normalized_settings,
                 **({"process_parameters": process_parameters_method} if process_parameters_method is not None else {}),
                 **({"process_blocks": process_blocks_method} if process_blocks_method is not None else {}),
                 "__module__": __name__,
@@ -409,6 +386,10 @@ def create_router_node(
                 function_name = get_function_name(function)
                 port = create_port_condition(function_name)
                 setattr(Ports, function_name, port)
+            elif isinstance(function, VellumIntegrationToolDefinition):
+                function_name = get_function_name(function)
+                port = create_port_condition(function_name)
+                setattr(Ports, function_name, port)
             elif isinstance(function, MCPServer):
                 tool_functions: List[MCPToolDefinition] = compile_mcp_tool_definition(function)
                 for tool_function in tool_functions:
@@ -483,6 +464,12 @@ def create_function_node(
             },
         )
         return node
+    elif isinstance(function, VellumIntegrationToolDefinition):
+        # TODO: Implement VellumIntegrationNode
+        raise NotImplementedError(
+            "VellumIntegrationToolDefinition support coming soon. "
+            "This will be implemented when the VellumIntegrationService is created."
+        )
     elif is_workflow_class(function):
         function.is_dynamic = True
         node = type(
@@ -572,5 +559,7 @@ def get_function_name(function: ToolBase) -> str:
     elif isinstance(function, ComposioToolDefinition):
         # model post init sets the name to the action if it's not set
         return function.name  # type: ignore[return-value]
+    elif isinstance(function, VellumIntegrationToolDefinition):
+        return function.name
     else:
         return snake_case(function.__name__)