vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

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

@@ -0,0 +1,633 @@
+import json
+import logging
+from typing import Any, Callable, Dict, Iterator, List, Optional, Type, Union, cast
+
+from pydash import snake_case
+
+from vellum import ChatMessage, PromptBlock
+from vellum.client.types.array_chat_message_content import ArrayChatMessageContent
+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.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.utils.json_encoder import VellumJsonEncoder
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.expressions.concat import ConcatExpression
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.integrations.composio_service import ComposioService
+from vellum.workflows.integrations.mcp_service import MCPService
+from vellum.workflows.integrations.vellum_integration_service import VellumIntegrationService
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.nodes.displayable.subworkflow_deployment_node.node import SubworkflowDeploymentNode
+from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
+from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.ports.port import Port
+from vellum.workflows.state import BaseState
+from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
+from vellum.workflows.types.definition import (
+    ComposioToolDefinition,
+    DeploymentDefinition,
+    MCPServer,
+    MCPToolDefinition,
+    Tool,
+    ToolBase,
+    VellumIntegrationToolDefinition,
+)
+from vellum.workflows.types.generics import is_workflow_class
+from vellum.workflows.utils.functions import compile_mcp_tool_definition, get_mcp_tool_name
+
+CHAT_HISTORY_VARIABLE = "chat_history"
+
+
+logger = logging.getLogger(__name__)
+
+
+class FunctionCallNodeMixin:
+    """Mixin providing common functionality for nodes that handle function calls."""
+
+    function_call_output: List[PromptOutput]
+
+    def _handle_tool_exception(self, e: Exception, tool_type: str, tool_name: str) -> None:
+        """
+        Re-raise exceptions with contextual information while preserving NodeException details.
+
+        Args:
+            e: The caught exception
+            tool_type: Type of tool (e.g., "function", "MCP tool", "Vellum Integration tool")
+            tool_name: Name of the tool that failed
+        """
+        if isinstance(e, NodeException):
+            # Preserve original error code and raw_data while adding context
+            raise NodeException(
+                message=f"Error executing {tool_type} '{tool_name}': {e.message}",
+                code=e.code,
+                raw_data=e.raw_data,
+            ) from e
+        else:
+            raise NodeException(
+                message=f"Error executing {tool_type} '{tool_name}': {str(e)}",
+                code=WorkflowErrorCode.NODE_EXECUTION,
+            ) from e
+
+    def _extract_function_arguments(self) -> dict:
+        """Extract arguments from function call output."""
+        current_index = getattr(self, "state").current_prompt_output_index
+        if self.function_call_output and len(self.function_call_output) > current_index:
+            function_call = self.function_call_output[current_index]
+            if function_call.type == "FUNCTION_CALL" and function_call.value is not None:
+                return function_call.value.arguments or {}
+        return {}
+
+    def _extract_function_call_id(self) -> Optional[str]:
+        """Extract function call ID from function call output."""
+        current_index = getattr(self, "state").current_prompt_output_index
+        if self.function_call_output and len(self.function_call_output) > current_index:
+            function_call = self.function_call_output[current_index]
+            if function_call.type == "FUNCTION_CALL" and function_call.value is not None:
+                return function_call.value.id
+        return None
+
+    def _add_function_result_to_chat_history(self, result: Any, state: ToolCallingState) -> None:
+        """Add function execution result to chat history."""
+        function_call_id = self._extract_function_call_id()
+        state.chat_history.append(
+            ChatMessage(
+                role="FUNCTION",
+                content=StringChatMessageContent(value=json.dumps(result, cls=VellumJsonEncoder)),
+                source=function_call_id,
+            )
+        )
+        with state.__quiet__():
+            state.current_function_calls_processed += 1
+            state.current_prompt_output_index += 1
+
+
+class ToolPromptNode(InlinePromptNode[ToolCallingState]):
+    max_prompt_iterations: Optional[int] = 25
+
+    class Trigger(InlinePromptNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ATTRIBUTES
+
+    def run(self) -> Iterator[BaseOutput]:
+        if self.max_prompt_iterations is not None and self.state.prompt_iterations >= self.max_prompt_iterations:
+            max_iterations_message = f"Maximum number of prompt iterations `{self.max_prompt_iterations}` reached."
+            raise NodeException(message=max_iterations_message, code=WorkflowErrorCode.NODE_EXECUTION)
+
+        generator = super().run()
+        with self.state.__quiet__():
+            self.state.current_prompt_output_index = 0
+            self.state.current_function_calls_processed = 0
+            self.state.prompt_iterations += 1
+        for output in generator:
+            if output.name == InlinePromptNode.Outputs.results.name and output.value:
+                prompt_outputs = cast(List[PromptOutput], output.value)
+                chat_contents: List[ArrayChatMessageContentItem] = []
+                for prompt_output in prompt_outputs:
+                    if prompt_output.type == "STRING":
+                        chat_contents.append(StringChatMessageContent(value=prompt_output.value))
+                    elif prompt_output.type == "FUNCTION_CALL" and prompt_output.value:
+                        raw_function_call = prompt_output.value.model_dump()
+                        if "state" in raw_function_call:
+                            del raw_function_call["state"]
+                        chat_contents.append(
+                            FunctionCallChatMessageContent(
+                                value=FunctionCallChatMessageContentValue.model_validate(raw_function_call)
+                            )
+                        )
+
+                if len(chat_contents) == 1:
+                    if chat_contents[0].type == "STRING":
+                        self.state.chat_history.append(ChatMessage(role="ASSISTANT", text=chat_contents[0].value))
+                    else:
+                        self.state.chat_history.append(ChatMessage(role="ASSISTANT", content=chat_contents[0]))
+                else:
+                    self.state.chat_history.append(
+                        ChatMessage(role="ASSISTANT", content=ArrayChatMessageContent(value=chat_contents))
+                    )
+
+            yield output
+
+
+class RouterNode(BaseNode[ToolCallingState]):
+    """Router node that handles routing to function nodes based on outputs."""
+
+    __exclude_from_monitoring__: bool = True  # Exclude from monitoring views
+
+
+class DynamicSubworkflowDeploymentNode(SubworkflowDeploymentNode[ToolCallingState], FunctionCallNodeMixin):
+    """Node that executes a deployment definition with function call output."""
+
+    def run(self) -> Iterator[BaseOutput]:
+        arguments = self._extract_function_arguments()
+
+        # Mypy doesn't like instance assignments of class attributes. It's safe in our case tho bc it's what
+        # we do in the `__init__` method. Long term, instead of the function_call_output attribute above, we
+        # want to do:
+        # ```python
+        # subworkflow_inputs = tool_prompt_node.Outputs.results[0]['value']['arguments'].if_(
+        #     tool_prompt_node.Outputs.results[0]['type'].equals('FUNCTION_CALL'),
+        #     {},
+        # )
+        # ```
+        self.subworkflow_inputs = arguments  # type:ignore[misc]
+
+        # Call the parent run method to execute the subworkflow
+        outputs = {}
+        for output in super().run():
+            if output.is_fulfilled:
+                outputs[output.name] = output.value
+            yield output
+
+        # Add the result to the chat history
+        self._add_function_result_to_chat_history(outputs, self.state)
+
+
+class DynamicInlineSubworkflowNode(
+    InlineSubworkflowNode[ToolCallingState, BaseInputs, BaseState], FunctionCallNodeMixin
+):
+    """Node that executes an inline subworkflow with function call output."""
+
+    def run(self) -> Iterator[BaseOutput]:
+        arguments = self._extract_function_arguments()
+
+        self.subworkflow_inputs = arguments  # type: ignore[misc]
+
+        # Call the parent run method to execute the subworkflow with proper streaming
+        outputs = {}
+
+        for output in super().run():
+            if output.is_fulfilled:
+                outputs[output.name] = output.value
+            yield output
+
+        # Add the result to the chat history
+        self._add_function_result_to_chat_history(outputs, self.state)
+
+
+class FunctionNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
+    """Node that executes a regular Python function with function call output."""
+
+    function_definition: Callable[..., Any]
+
+    class Outputs(BaseNode.Outputs):
+        result: Any
+
+    def run(self) -> Iterator[BaseOutput]:
+        arguments = self._extract_function_arguments()
+
+        try:
+            result = self.function_definition(**arguments)
+        except Exception as e:
+            self._handle_tool_exception(e, "function", self.function_definition.__name__)
+
+        # Add the result to the chat history
+        self._add_function_result_to_chat_history(result, self.state)
+
+        yield BaseOutput(name="result", value=result)
+
+
+class ComposioNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
+    """Node that executes a Composio tool with function call output."""
+
+    composio_tool: ComposioToolDefinition
+
+    def run(self) -> Iterator[BaseOutput]:
+        # Extract arguments from function call
+        arguments = self._extract_function_arguments()
+
+        try:
+            # Execute using ComposioService
+            composio_service = ComposioService()
+            if self.composio_tool.user_id is not None:
+                result = composio_service.execute_tool(
+                    tool_name=self.composio_tool.action, arguments=arguments, user_id=self.composio_tool.user_id
+                )
+            else:
+                result = composio_service.execute_tool(tool_name=self.composio_tool.action, arguments=arguments)
+        except Exception as e:
+            self._handle_tool_exception(e, "Composio tool", self.composio_tool.action)
+
+        # Add result to chat history
+        self._add_function_result_to_chat_history(result, self.state)
+
+        yield from []
+
+
+class MCPNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
+    """Node that executes an MCP tool with function call output."""
+
+    mcp_tool: MCPToolDefinition
+
+    class Outputs(BaseNode.Outputs):
+        result: Any
+
+    def run(self) -> Iterator[BaseOutput]:
+        arguments = self._extract_function_arguments()
+
+        try:
+            mcp_service = MCPService()
+            result = mcp_service.execute_tool(tool_def=self.mcp_tool, arguments=arguments)
+        except Exception as e:
+            self._handle_tool_exception(e, "MCP tool", self.mcp_tool.name)
+
+        # Add result to chat history
+        self._add_function_result_to_chat_history(result, self.state)
+
+        yield BaseOutput(name="result", value=result)
+
+
+class VellumIntegrationNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
+    """Node that executes a Vellum Integration tool with function call output."""
+
+    vellum_integration_tool: VellumIntegrationToolDefinition
+
+    class Outputs(BaseNode.Outputs):
+        result: Any
+
+    def run(self) -> Iterator[BaseOutput]:
+        arguments = self._extract_function_arguments()
+        vellum_client = self._context.vellum_client
+
+        try:
+            vellum_service = VellumIntegrationService(vellum_client)
+            result = vellum_service.execute_tool(
+                integration=self.vellum_integration_tool.integration_name,
+                provider=self.vellum_integration_tool.provider.value,
+                tool_name=self.vellum_integration_tool.name,
+                arguments=arguments,
+            )
+        except NodeException as e:
+            error_payload = {
+                "error": {
+                    "code": e.code.value,
+                    "message": e.message,
+                }
+            }
+            if e.raw_data is not None:
+                error_payload["error"]["raw_data"] = e.raw_data
+
+            self._add_function_result_to_chat_history(error_payload, self.state)
+            yield from []
+            return
+        except Exception as e:
+            self._handle_tool_exception(e, "Vellum Integration tool", self.vellum_integration_tool.name)
+
+        # Add result to chat history
+        self._add_function_result_to_chat_history(result, self.state)
+
+        yield BaseOutput(name="result", value=result)
+
+
+class ElseNode(BaseNode[ToolCallingState]):
+    """Node that executes when no function conditions match."""
+
+    __exclude_from_monitoring__: bool = True  # Exclude from monitoring views
+
+    class Ports(BaseNode.Ports):
+        # Redefined in the create_else_node function, but defined here to resolve mypy errors
+        loop_to_router = Port.on_if(
+            ToolCallingState.current_prompt_output_index.less_than(ToolPromptNode.Outputs.results.length())
+        )
+        loop_to_prompt = Port.on_elif(ToolCallingState.current_function_calls_processed.greater_than(0))
+        end = Port.on_else()
+
+    def run(self) -> BaseNode.Outputs:
+        with self.state.__quiet__():
+            self.state.current_prompt_output_index += 1
+        return self.Outputs()
+
+
+def create_tool_prompt_node(
+    ml_model: str,
+    blocks: List[Union[PromptBlock, Dict[str, Any]]],
+    functions: List[Tool],
+    prompt_inputs: Optional[EntityInputsInterface],
+    parameters: PromptParameters,
+    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[Tool] = functions
+    else:
+        prompt_functions = []
+
+    # Add a chat history block to blocks only if one doesn't already exist
+    has_chat_history_block = any(
+        (
+            (block["block_type"] if isinstance(block, dict) else block.block_type) == "VARIABLE"
+            and (
+                block["input_variable"]
+                if isinstance(block, dict)
+                else block.input_variable if isinstance(block, VariablePromptBlock) else None
+            )
+            == CHAT_HISTORY_VARIABLE
+        )
+        for block in blocks
+    )
+
+    if not has_chat_history_block:
+        blocks.append(
+            VariablePromptBlock(
+                block_type="VARIABLE",
+                input_variable=CHAT_HISTORY_VARIABLE,
+                state=None,
+                cache_config=None,
+            )
+        )
+
+    node_prompt_inputs = {
+        **(prompt_inputs or {}),
+        CHAT_HISTORY_VARIABLE: ConcatExpression[List[ChatMessage], List[ChatMessage]](
+            lhs=(prompt_inputs or {}).get(CHAT_HISTORY_VARIABLE, []),
+            rhs=ToolCallingState.chat_history,
+        ),
+    }
+
+    # 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(
+            "ToolPromptNode",
+            (ToolPromptNode,),
+            {
+                "ml_model": ml_model,
+                "blocks": blocks,
+                "functions": prompt_functions,  # Use converted functions for prompt layer
+                "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__,
+            },
+        ),
+    )
+    return node
+
+
+def create_router_node(
+    functions: List[Tool],
+    tool_prompt_node: Type[InlinePromptNode[ToolCallingState]],
+) -> Type[RouterNode]:
+    """Create a RouterNode with dynamic ports that route based on tool_prompt_node outputs."""
+
+    if functions and len(functions) > 0:
+        # Create dynamic ports and convert functions in a single loop
+        Ports = type("Ports", (), {})
+
+        # Collect all tool names
+        tool_names: List[str] = []
+        for function in functions:
+            if isinstance(function, MCPServer):
+                tool_functions: List[MCPToolDefinition] = compile_mcp_tool_definition(function)
+                for tool_function in tool_functions:
+                    tool_names.append(get_mcp_tool_name(tool_function))
+            else:
+                tool_names.append(get_function_name(function))
+
+        # Build conditions for each tool name
+        conditions = [
+            (
+                name,
+                ToolCallingState.current_prompt_output_index.less_than(tool_prompt_node.Outputs.results.length())
+                & tool_prompt_node.Outputs.results[ToolCallingState.current_prompt_output_index]["type"].equals(
+                    "FUNCTION_CALL"
+                )
+                & tool_prompt_node.Outputs.results[ToolCallingState.current_prompt_output_index]["value"][
+                    "name"
+                ].equals(name),
+            )
+            for name in tool_names
+        ]
+
+        # Assign ports: first condition uses on_if, subsequent ones use on_elif
+        for idx, (name, condition) in enumerate(conditions):
+            port = Port.on_if(condition) if idx == 0 else Port.on_elif(condition)
+            setattr(Ports, name, port)
+
+        # Add the else port for when no function conditions match
+        setattr(Ports, "default", Port.on_else())
+    else:
+        # If no functions exist, create a simple Ports class with just a default port
+        Ports = type("Ports", (), {"default": Port(default=True)})
+
+    node = cast(
+        Type[RouterNode],
+        type(
+            "RouterNode",
+            (RouterNode,),
+            {
+                "Ports": Ports,
+                "prompt_outputs": tool_prompt_node.Outputs.results,
+                "__module__": __name__,
+            },
+        ),
+    )
+    return node
+
+
+def create_function_node(
+    function: ToolBase,
+    tool_prompt_node: Type[ToolPromptNode],
+) -> Type[BaseNode]:
+    """
+    Create a FunctionNode class for a given function.
+
+    For workflow functions: BaseNode
+    For regular functions: BaseNode with direct function call
+
+    Args:
+        function: The function to create a node for
+        tool_prompt_node: The tool prompt node class
+    """
+    if isinstance(function, DeploymentDefinition):
+        deployment = function.deployment_id or function.deployment_name
+        release_tag = function.release_tag
+
+        node = type(
+            f"DynamicSubworkflowDeploymentNode_{deployment}",
+            (DynamicSubworkflowDeploymentNode,),
+            {
+                "deployment": deployment,
+                "release_tag": release_tag,
+                "function_call_output": tool_prompt_node.Outputs.results,
+                "__module__": __name__,
+            },
+        )
+
+        return node
+
+    elif isinstance(function, ComposioToolDefinition):
+        node = type(
+            f"ComposioNode_{function.name}",
+            (ComposioNode,),
+            {
+                "composio_tool": function,
+                "function_call_output": tool_prompt_node.Outputs.results,
+                "__module__": __name__,
+            },
+        )
+        return node
+    elif isinstance(function, VellumIntegrationToolDefinition):
+        node = type(
+            f"VellumIntegrationNode_{function.name}",
+            (VellumIntegrationNode,),
+            {
+                "vellum_integration_tool": function,
+                "function_call_output": tool_prompt_node.Outputs.results,
+                "__module__": __name__,
+            },
+        )
+        return node
+    elif is_workflow_class(function):
+        function.is_dynamic = True
+        node = type(
+            f"DynamicInlineSubworkflowNode_{function.__name__}",
+            (DynamicInlineSubworkflowNode,),
+            {
+                "subworkflow": function,
+                "function_call_output": tool_prompt_node.Outputs.results,
+                "__module__": __name__,
+            },
+        )
+    else:
+
+        def create_function_wrapper(func):
+            def wrapper(self, **kwargs):
+                merged_kwargs = kwargs.copy()
+                inputs = getattr(func, "__vellum_inputs__", {})
+                if inputs:
+                    for param_name, param_ref in inputs.items():
+                        if isinstance(param_ref, BaseDescriptor):
+                            resolved_value = param_ref.resolve(self.state)
+                        else:
+                            resolved_value = param_ref
+                        merged_kwargs[param_name] = resolved_value
+
+                return func(**merged_kwargs)
+
+            # Preserve the original function name for better error messages
+            wrapper.__name__ = func.__name__
+            return wrapper
+
+        node = type(
+            f"FunctionNode_{function.__name__}",
+            (FunctionNode,),
+            {
+                "function_definition": create_function_wrapper(function),
+                "function_call_output": tool_prompt_node.Outputs.results,
+                "__module__": __name__,
+            },
+        )
+
+    return node
+
+
+def create_mcp_tool_node(
+    tool_def: MCPToolDefinition,
+    tool_prompt_node: Type[ToolPromptNode],
+) -> Type[BaseNode]:
+    node = type(
+        f"MCPNode_{tool_def.name}",
+        (MCPNode,),
+        {
+            "mcp_tool": tool_def,
+            "function_call_output": tool_prompt_node.Outputs.results,
+            "__module__": __name__,
+        },
+    )
+    return node
+
+
+def create_else_node(
+    tool_prompt_node: Type[ToolPromptNode],
+) -> Type[ElseNode]:
+    class Ports(ElseNode.Ports):
+        loop_to_router = Port.on_if(
+            ToolCallingState.current_prompt_output_index.less_than(tool_prompt_node.Outputs.results.length())
+        )
+        loop_to_prompt = Port.on_elif(ToolCallingState.current_function_calls_processed.greater_than(0))
+        end = Port.on_else()
+
+    node = cast(
+        Type[ElseNode],
+        type(
+            f"{tool_prompt_node.__name__}_ElseNode",
+            (ElseNode,),
+            {
+                "Ports": Ports,
+                "__module__": __name__,
+            },
+        ),
+    )
+    return node
+
+
+def get_function_name(function: ToolBase) -> str:
+    if isinstance(function, DeploymentDefinition):
+        name = str(function.deployment_id or function.deployment_name)
+        return name.replace("-", "")
+    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__)

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

@@ -0,0 +1,3 @@
+from .node import WebSearchNode
+
+__all__ = ["WebSearchNode"]