vellum-ai 1.0.10__py3-none-any.whl → 1.1.0__py3-none-any.whl

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

@@ -1,6 +1,6 @@
 import json
 from uuid import uuid4
-from typing import Any, Iterator, List
+from typing import Any, Iterator
 
 from vellum import ChatMessage
 from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
@@ -15,6 +15,7 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
+from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
 from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_function_node, create_tool_router_node
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState, StateMeta
@@ -44,7 +45,7 @@ def test_port_condition_match_function_name():
     )
 
     # AND a state with a function call to the first function
-    state = BaseState(
+    state = ToolCallingState(
         meta=StateMeta(
             node_outputs={
                 router_node.Outputs.results: [
@@ -116,12 +117,6 @@ def test_tool_calling_node_inline_workflow_context():
     )
     function_node._context = parent_context
 
-    # Create a state with chat_history for the function node
-    class TestState(BaseState):
-        chat_history: List[ChatMessage] = []
-
-    function_node.state = TestState(meta=StateMeta(node_outputs={tool_router_node.Outputs.text: '{"arguments": {}}'}))
-
     # WHEN the function node runs
     outputs = list(function_node.run())
 

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

@@ -5,6 +5,8 @@ from typing import Any, Callable, Dict, Iterator, List, Optional, Type, Union, c
 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.function_definition import FunctionDefinition
@@ -59,6 +61,9 @@ class FunctionCallNodeMixin:
                 content=StringChatMessageContent(value=json.dumps(result, cls=DefaultStateEncoder)),
             )
         )
+        with state.__quiet__():
+            state.current_function_calls_processed += 1
+            state.current_prompt_output_index += 1
 
 
 class ToolRouterNode(InlinePromptNode[ToolCallingState]):
@@ -73,32 +78,52 @@ class ToolRouterNode(InlinePromptNode[ToolCallingState]):
             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 == "results" and output.value:
-                values = cast(List[Any], output.value)
-                if values and len(values) > 0:
-                    if values[0].type == "STRING":
-                        self.state.chat_history.append(ChatMessage(role="ASSISTANT", text=values[0].value))
-                    elif values[0].type == "FUNCTION_CALL":
-                        self.state.prompt_iterations += 1
-
-                        function_call = values[0].value
-                        if function_call is not None:
-                            self.state.chat_history.append(
-                                ChatMessage(
-                                    role="ASSISTANT",
-                                    content=FunctionCallChatMessageContent(
-                                        value=FunctionCallChatMessageContentValue(
-                                            name=function_call.name,
-                                            arguments=function_call.arguments,
-                                            id=function_call.id,
-                                        ),
-                                    ),
-                                )
+            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."""
+
+    class Trigger(BaseNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ATTRIBUTES
+
+    def run(self) -> Iterator[BaseOutput]:
+        # Router node doesn't process outputs or create chat messages
+        # It just handles the routing logic via its ports
+        yield from []
+
+
 class DynamicSubworkflowDeploymentNode(SubworkflowDeploymentNode[ToolCallingState], FunctionCallNodeMixin):
     """Node that executes a deployment definition with function call output."""
 
@@ -225,42 +250,51 @@ class MCPNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
         yield from []
 
 
-def _hydrate_composio_tool_definition(tool_def: ComposioToolDefinition) -> ComposioToolDefinition:
+class ElseNode(BaseNode[ToolCallingState]):
+    """Node that executes when no function conditions match."""
+
+    class Ports(BaseNode.Ports):
+        # Redefined in the create_else_node function, but defined here to resolve mypy errors
+        loop = Port.on_if(
+            ToolCallingState.current_prompt_output_index.less_than(1)
+            | 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 _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:
-        ComposioToolDefinition with detailed parameters and description
+        FunctionDefinition with detailed parameters and description
     """
     try:
         composio_service = ComposioService()
         tool_details = composio_service.get_tool_by_slug(tool_def.action)
 
-        # Extract toolkit information from API response
-        toolkit_info = tool_details.get("toolkit", {})
-        toolkit_slug = (
-            toolkit_info.get("slug", tool_def.toolkit) if isinstance(toolkit_info, dict) else tool_def.toolkit
-        )
-
-        # Create a version of the tool definition with proper field extraction
-        return ComposioToolDefinition(
-            type=tool_def.type,
-            toolkit=toolkit_slug.upper() if toolkit_slug else tool_def.toolkit,
-            action=tool_details.get("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),
-            display_name=tool_details.get("name", tool_def.display_name),
-            parameters=tool_details.get("input_parameters", tool_def.parameters),
-            version=tool_details.get("version", tool_def.version),
-            tags=tool_details.get("tags", tool_def.tags),
-            user_id=tool_def.user_id,
+            parameters=tool_details.get("input_parameters", {}),
         )
 
     except Exception as e:
-        # If hydration fails (including no API key), log and return original
+        # 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 tool_def
+        return FunctionDefinition(
+            name=tool_def.name,
+            description=tool_def.description,
+            parameters={},
+        )
 
 
 def hydrate_mcp_tool_definitions(server_def: MCPServer) -> List[MCPToolDefinition]:
@@ -303,8 +337,13 @@ def create_tool_router_node(
             return Port.on_if(
                 LazyReference(
                     lambda: (
-                        node.Outputs.results[0]["type"].equals("FUNCTION_CALL")
-                        & node.Outputs.results[0]["value"]["name"].equals(fn_name)
+                        ToolCallingState.current_prompt_output_index.less_than(node.Outputs.results.length())
+                        & node.Outputs.results[ToolCallingState.current_prompt_output_index]["type"].equals(
+                            "FUNCTION_CALL"
+                        )
+                        & node.Outputs.results[ToolCallingState.current_prompt_output_index]["value"]["name"].equals(
+                            fn_name
+                        )
                     )
                 )
             )
@@ -313,13 +352,7 @@ def create_tool_router_node(
             if isinstance(function, ComposioToolDefinition):
                 # Get Composio tool details and hydrate the function definition
                 enhanced_function = _hydrate_composio_tool_definition(function)
-                prompt_functions.append(
-                    FunctionDefinition(
-                        name=enhanced_function.name,
-                        description=enhanced_function.description,
-                        parameters=enhanced_function.parameters,
-                    )
-                )
+                prompt_functions.append(enhanced_function)
                 # Create port for this function (using original function for get_function_name)
                 function_name = get_function_name(function)
                 port = create_port_condition(function_name)
@@ -402,6 +435,69 @@ def create_tool_router_node(
     return node
 
 
+def create_router_node(
+    functions: List[Tool],
+    tool_router_node: Type[ToolRouterNode],
+) -> Type[RouterNode]:
+    """Create a RouterNode with the same ports as ToolRouterNode."""
+
+    if functions and len(functions) > 0:
+        # Create dynamic ports and convert functions in a single loop
+        Ports = type("Ports", (), {})
+
+        def create_port_condition(fn_name):
+            return Port.on_if(
+                LazyReference(
+                    lambda: (
+                        ToolCallingState.current_prompt_output_index.less_than(
+                            tool_router_node.Outputs.results.length()
+                        )
+                        & tool_router_node.Outputs.results[ToolCallingState.current_prompt_output_index]["type"].equals(
+                            "FUNCTION_CALL"
+                        )
+                        & tool_router_node.Outputs.results[ToolCallingState.current_prompt_output_index]["value"][
+                            "name"
+                        ].equals(fn_name)
+                    )
+                )
+            )
+
+        for function in functions:
+            if isinstance(function, ComposioToolDefinition):
+                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] = hydrate_mcp_tool_definitions(function)
+                for tool_function in tool_functions:
+                    name = get_mcp_tool_name(tool_function)
+                    port = create_port_condition(name)
+                    setattr(Ports, name, port)
+            else:
+                function_name = get_function_name(function)
+                port = create_port_condition(function_name)
+                setattr(Ports, function_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,
+                "__module__": __name__,
+            },
+        ),
+    )
+    return node
+
+
 def create_function_node(
     function: ToolBase,
     tool_router_node: Type[ToolRouterNode],
@@ -485,6 +581,27 @@ def create_mcp_tool_node(
     return node
 
 
+def create_else_node(
+    tool_router_node: Type[ToolRouterNode],
+) -> Type[ElseNode]:
+    class Ports(ElseNode.Ports):
+        loop = Port.on_if(
+            ToolCallingState.current_prompt_output_index.less_than(tool_router_node.Outputs.results.length())
+            | ToolCallingState.current_function_calls_processed.greater_than(0)
+        )
+        end = Port.on_else()
+
+    node = type(
+        f"{tool_router_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)

vellum/workflows/state/context.py

@@ -1,9 +1,11 @@
 from functools import cached_property
 from queue import Queue
+from uuid import uuid4
 from typing import TYPE_CHECKING, Dict, List, Optional, Type
 
 from vellum import Vellum
 from vellum.workflows.context import ExecutionContext, get_execution_context
+from vellum.workflows.events.types import ExternalParentContext
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
@@ -25,8 +27,17 @@ class WorkflowContext:
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
         self._node_output_mocks_map: Dict[Type[BaseOutputs], List[MockNodeExecution]] = {}
         self._execution_context = get_execution_context()
-        if not self._execution_context.parent_context and execution_context:
-            self._execution_context = execution_context
+
+        if execution_context is not None:
+
+            self._execution_context.trace_id = execution_context.trace_id
+
+            if execution_context.parent_context is not None:
+                self._execution_context.parent_context = execution_context.parent_context
+
+        if self._execution_context.parent_context is None:
+            self._execution_context.parent_context = ExternalParentContext(span_id=uuid4())
+
         self._generated_files = generated_files
 
     @cached_property

vellum/workflows/types/definition.py

@@ -2,13 +2,12 @@ import importlib
 import inspect
 from types import FrameType
 from uuid import UUID
-from typing import Annotated, Any, Dict, List, Literal, Optional, Union
+from typing import Annotated, Any, Dict, Literal, Optional, Union
 
 from pydantic import BeforeValidator
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.code_resource_definition import CodeResourceDefinition as ClientCodeResourceDefinition
-from vellum.client.types.vellum_secret import VellumSecret
 from vellum.workflows.constants import AuthorizationType
 from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 
@@ -111,11 +110,6 @@ class ComposioToolDefinition(UniversalBaseModel):
     toolkit: str  # "GITHUB", "SLACK", etc.
     action: str  # Specific action like "GITHUB_CREATE_AN_ISSUE"
     description: str
-
-    display_name: Optional[str] = None
-    parameters: Optional[Dict[str, Any]] = None
-    version: Optional[str] = None
-    tags: Optional[List[str]] = None
     user_id: Optional[str] = None
 
     @property
@@ -125,12 +119,13 @@ class ComposioToolDefinition(UniversalBaseModel):
 
 
 class MCPServer(UniversalBaseModel):
+    type: Literal["MCP_SERVER"] = "MCP_SERVER"
     name: str
     url: str
     authorization_type: AuthorizationType = AuthorizationType.BEARER_TOKEN
     bearer_token_value: Optional[Union[str, EnvironmentVariableReference]] = None
     api_key_header_key: Optional[str] = None
-    api_key_header_value: Optional[Union[str, VellumSecret]] = None
+    api_key_header_value: Optional[Union[str, EnvironmentVariableReference]] = None
 
     model_config = {"arbitrary_types_allowed": True}
 

vellum/workflows/types/tests/test_definition.py

@@ -1,7 +1,6 @@
 import pytest
 from uuid import UUID
 
-from vellum.client.types.vellum_secret import VellumSecret
 from vellum.workflows.constants import AuthorizationType
 from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition, MCPServer, MCPToolDefinition
@@ -48,7 +47,7 @@ def test_composio_tool_definition_creation():
     assert composio_tool.toolkit == "GITHUB"
     assert composio_tool.action == "GITHUB_CREATE_AN_ISSUE"
     assert composio_tool.description == "Create a new issue in a GitHub repository"
-    assert composio_tool.display_name is None
+    assert composio_tool.user_id is None
     assert composio_tool.name == "github_create_an_issue"
 
 
@@ -98,7 +97,7 @@ def test_mcp_tool_definition_creation_api_key():
             url="https://api.githubcopilot.com/mcp/",
             authorization_type=AuthorizationType.API_KEY,
             api_key_header_key="Authorization",
-            api_key_header_value=VellumSecret(name="GITHUB_PERSONAL_ACCESS_TOKEN"),
+            api_key_header_value=EnvironmentVariableReference(name="GITHUB_PERSONAL_ACCESS_TOKEN"),
         ),
         parameters={
             "type": "object",
@@ -115,7 +114,7 @@ def test_mcp_tool_definition_creation_api_key():
     assert mcp_tool.server.url == "https://api.githubcopilot.com/mcp/"
     assert mcp_tool.server.authorization_type == AuthorizationType.API_KEY
     assert mcp_tool.server.api_key_header_key == "Authorization"
-    assert isinstance(mcp_tool.server.api_key_header_value, VellumSecret)
+    assert isinstance(mcp_tool.server.api_key_header_value, EnvironmentVariableReference)
     assert mcp_tool.server.api_key_header_value.name == "GITHUB_PERSONAL_ACCESS_TOKEN"
     assert mcp_tool.parameters == {
         "type": "object",

vellum/workflows/utils/functions.py

@@ -221,7 +221,7 @@ def compile_workflow_deployment_function_definition(
     deployment = deployment_config["deployment"]
     release_tag = deployment_config["release_tag"]
 
-    workflow_deployment_release = vellum_client.release_reviews.retrieve_workflow_deployment_release(
+    workflow_deployment_release = vellum_client.workflow_deployments.retrieve_workflow_deployment_release(
         deployment, release_tag
     )
 

vellum/workflows/utils/tests/test_functions.py

@@ -448,7 +448,7 @@ def test_compile_workflow_deployment_function_definition__just_name():
     mock_release = Mock()
     mock_release.workflow_version.input_variables = []
     mock_release.description = "This is a test deployment"
-    mock_client.release_reviews.retrieve_workflow_deployment_release.return_value = mock_release
+    mock_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = mock_release
 
     deployment_config = {"deployment": "my_deployment", "release_tag": "latest"}
 
@@ -494,7 +494,7 @@ def test_compile_workflow_deployment_function_definition__all_args():
 
     mock_release.workflow_version.input_variables = mock_inputs
     mock_release.description = "This is a test deployment"
-    mock_client.release_reviews.retrieve_workflow_deployment_release.return_value = mock_release
+    mock_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = mock_release
 
     deployment_config = {"deployment": "my_deployment", "release_tag": "latest"}
 
@@ -562,7 +562,7 @@ def test_compile_workflow_deployment_function_definition__defaults():
 
     mock_release.workflow_version.input_variables = mock_inputs
     mock_release.description = "This is a test deployment"
-    mock_client.release_reviews.retrieve_workflow_deployment_release.return_value = mock_release
+    mock_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = mock_release
 
     deployment_config = {"deployment": "my_deployment", "release_tag": "latest"}
 

{vellum_ai-1.0.10.dist-info → vellum_ai-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.0.10
+Version: 1.1.0
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0