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

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

@@ -0,0 +1,136 @@
+import logging
+from typing import Any, ClassVar, Dict, List, Optional
+
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.core import ApiError
+from vellum.workflows.errors.types import WorkflowErrorCode
+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
+
+logger = logging.getLogger(__name__)
+
+
+class WebSearchResult(UniversalBaseModel):
+    title: str
+    url: str
+    snippet: Optional[str] = None
+
+
+class WebSearchNode(BaseNode[StateType]):
+    """
+    Used to perform web search using SerpAPI.
+
+    query: str - The search query to execute
+    """
+
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:magnifying-glass"
+        color = "lightBlue"
+
+    query: ClassVar[str] = ""
+
+    class Outputs(BaseOutputs):
+        """
+        The outputs of the WebSearchNode.
+
+        text: str - Concatenated search result snippets with titles
+        urls: List[str] - List of URLs from search results
+        results: List[Dict[str, Any]] - Raw search results from SerpAPI
+        """
+
+        text: str
+        urls: List[str]
+        results: List[Dict[str, Any]]
+
+    def _validate_query(self) -> str:
+        if not self.query or not isinstance(self.query, str) or not self.query.strip():
+            raise NodeException(
+                "Query is required and must be a non-empty string", code=WorkflowErrorCode.INVALID_INPUTS
+            )
+
+        return self.query.strip()
+
+    def run(self) -> Outputs:
+        """Run the WebSearchNode to perform web search via SerpAPI."""
+        query = self._validate_query()
+
+        try:
+            response = self._context.vellum_client.integrations.execute_integration_tool(
+                integration_name="SERPAPI",
+                integration_provider="COMPOSIO",
+                tool_name="SERPAPI_SEARCH",
+                arguments={"query": query},
+            )
+        except ApiError as e:
+            if e.status_code == 403 and isinstance(e.body, dict):
+                raise NodeException(
+                    message=e.body.get("detail", "You must provide authenticate with the search provider."),
+                    code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+                ) from e
+            elif e.status_code == 404 and isinstance(e.body, dict):
+                error_message = e.body.get("detail", "Tool not found")
+                raise NodeException(message=error_message, code=WorkflowErrorCode.PROVIDER_ERROR) from e
+
+            error_message = (
+                e.body.get("detail", f"Failed to execute SerpAPI search: {str(e)}")
+                if isinstance(e.body, dict)
+                else f"Failed to execute SerpAPI search: {str(e)}"
+            )
+            raise NodeException(message=error_message, code=WorkflowErrorCode.INTERNAL_ERROR) from e
+
+        web_search_results = self._parse_serpapi_results(response.data)
+        return self._web_search_results_to_outputs(web_search_results)
+
+    def _parse_serpapi_results(self, data: Dict) -> List[WebSearchResult]:
+        if "error" in data:
+            error_msg = data["error"]
+            logger.error(f"SerpAPI returned error: {error_msg}")
+            raise NodeException(f"SerpAPI error: {error_msg}", code=WorkflowErrorCode.PROVIDER_ERROR)
+
+        try:
+            results: Dict = data["results"]
+            organic_results: List[Dict] = results["organic_results"]
+        except KeyError as e:
+            raise NodeException("Unable to parse web search response", code=WorkflowErrorCode.PROVIDER_ERROR) from e
+
+        web_search_results: List[WebSearchResult] = []
+        for result in organic_results:
+            title = result["title"]
+            url = result["link"]
+            snippet = result.get("snippet")
+
+            web_search_results.append(
+                WebSearchResult(
+                    title=title,
+                    url=url,
+                    snippet=snippet,
+                )
+            )
+
+        return web_search_results
+
+    def _web_search_results_to_outputs(self, web_search_results: List[WebSearchResult]) -> Outputs:
+        text: str = ""
+        urls: List[str] = []
+        results: List[Dict] = []
+
+        for web_search_result in web_search_results:
+            if text:
+                text += "\n\n"
+
+            text += f"Title: {web_search_result.title}"
+            if web_search_result.snippet:
+                text += f"\nSnippet: {web_search_result.snippet}"
+            if web_search_result.url:
+                text += f"\nURL: {web_search_result.url}"
+                urls.append(web_search_result.url)
+
+            results.append(web_search_result.model_dump(mode="json"))
+
+        return self.Outputs(
+            text=text,
+            urls=urls,
+            results=results,
+        )

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

@@ -0,0 +1,246 @@
+import pytest
+from unittest.mock import MagicMock
+
+from vellum.client.core.api_error import ApiError
+from vellum.client.types.composio_execute_tool_response import ComposioExecuteToolResponse
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.state import BaseState
+from vellum.workflows.state.base import StateMeta
+
+from ..node import WebSearchNode
+
+
+@pytest.fixture
+def base_node_setup(vellum_client):
+    """Basic node setup with required inputs."""
+
+    class Inputs(BaseInputs):
+        query: str
+
+    class State(BaseState):
+        pass
+
+    class TestableWebSearchNode(WebSearchNode):
+        query = Inputs.query
+
+    state = State(meta=StateMeta(workflow_inputs=Inputs(query="test query")))
+    context = MagicMock()
+    context.vellum_client = vellum_client
+    node = TestableWebSearchNode(state=state, context=context)
+    return node
+
+
+def test_successful_search_with_results(base_node_setup):
+    """Test successful SerpAPI search with typical results."""
+
+    # GIVEN a mock Composio SerpAPI response with results
+    mock_response_data = {
+        "results": {
+            "organic_results": [
+                {
+                    "title": "First Result",
+                    "snippet": "This is the first search result snippet",
+                    "link": "https://example1.com",
+                    "position": 1,
+                },
+                {
+                    "title": "Second Result",
+                    "link": "https://example2.com",
+                    "position": 2,
+                },
+                {
+                    "title": "Third Result",
+                    "link": "https://example3.com",
+                    "position": 3,
+                },
+            ]
+        }
+    }
+
+    mock_response = ComposioExecuteToolResponse(provider="COMPOSIO", data=mock_response_data)
+    base_node_setup._context.vellum_client.integrations.execute_integration_tool = MagicMock(return_value=mock_response)
+
+    # WHEN we run the node
+    outputs = base_node_setup.run()
+
+    # THEN the text output should be properly formatted with labels
+    expected_text = """\
+Title: First Result
+Snippet: This is the first search result snippet
+URL: https://example1.com
+
+Title: Second Result
+URL: https://example2.com
+
+Title: Third Result
+URL: https://example3.com\
+"""
+    assert outputs.text == expected_text
+
+    # AND URLs should be extracted correctly
+    assert outputs.urls == ["https://example1.com", "https://example2.com", "https://example3.com"]
+
+    # AND raw results should be serialized WebSearchResult models
+    assert len(outputs.results) == 3
+    assert outputs.results == [
+        {
+            "title": "First Result",
+            "snippet": "This is the first search result snippet",
+            "url": "https://example1.com",
+        },
+        {
+            "title": "Second Result",
+            "snippet": None,
+            "url": "https://example2.com",
+        },
+        {
+            "title": "Third Result",
+            "snippet": None,
+            "url": "https://example3.com",
+        },
+    ]
+
+    # AND the Vellum client should have been called with the correct parameters
+    base_node_setup._context.vellum_client.integrations.execute_integration_tool.assert_called_once_with(
+        integration_name="SERPAPI",
+        integration_provider="COMPOSIO",
+        tool_name="SERPAPI_SEARCH",
+        arguments={"query": "test query"},
+    )
+
+
+def test_authentication_error_401(base_node_setup):
+    """Test authentication error raises NodeException with PROVIDER_CREDENTIALS_UNAVAILABLE."""
+
+    base_node_setup._context.vellum_client.integrations.execute_integration_tool = MagicMock(
+        side_effect=ApiError(status_code=403, body={"detail": "You must authenticate with the search provider."})
+    )
+
+    # WHEN we run the node
+    with pytest.raises(NodeException) as exc_info:
+        base_node_setup.run()
+
+    # THEN it should raise the appropriate error with correct code
+    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE
+    assert exc_info.value.message == "You must authenticate with the search provider."
+
+
+def test_serpapi_error_in_response(base_node_setup):
+    """Test SerpAPI error field in response raises NodeException."""
+
+    # GIVEN SerpAPI returns an error in the response data
+    mock_response = ComposioExecuteToolResponse(provider="COMPOSIO", data={"error": "Invalid search parameters"})
+    base_node_setup._context.vellum_client.integrations.execute_integration_tool = MagicMock(return_value=mock_response)
+
+    # WHEN we run the node
+    with pytest.raises(NodeException) as exc_info:
+        base_node_setup.run()
+
+    # THEN it should raise the appropriate error
+    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_ERROR
+    assert exc_info.value.message == "SerpAPI error: Invalid search parameters"
+
+
+def test_empty_query_validation(vellum_client):
+    """Test empty query raises validation error."""
+
+    # GIVEN a node with an empty query
+    class TestNode(WebSearchNode):
+        query = ""
+
+    context = MagicMock()
+    context.vellum_client = vellum_client
+    node = TestNode(state=BaseState(meta=StateMeta(workflow_inputs=BaseInputs())), context=context)
+
+    # WHEN we run the node
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN it should raise a validation error
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert exc_info.value.message == "Query is required and must be a non-empty string"
+
+
+def test_empty_organic_results(base_node_setup):
+    """Test handling of empty search results."""
+    # GIVEN Composio SerpAPI returns no organic results in new format
+    mock_response = ComposioExecuteToolResponse(provider="COMPOSIO", data={"results": {"organic_results": []}})
+    base_node_setup._context.vellum_client.integrations.execute_integration_tool = MagicMock(return_value=mock_response)
+
+    # WHEN we run the node
+    outputs = base_node_setup.run()
+
+    # THEN all outputs should be empty
+    assert outputs.text == ""
+    assert outputs.urls == []
+    assert outputs.results == []
+
+
+def test_missing_fields_in_results(base_node_setup):
+    """Test handling of results with missing optional snippet field."""
+    # GIVEN Composio SerpAPI returns results with some missing snippet fields
+    mock_response_data = {
+        "results": {
+            "organic_results": [
+                {
+                    "title": "Result with snippet",
+                    "link": "https://example1.com",
+                    "snippet": "This has a snippet",
+                },
+                {
+                    "title": "Result without snippet",
+                    "link": "https://example2.com",
+                },
+                {
+                    "title": "Another with snippet",
+                    "link": "https://example3.com",
+                    "snippet": "Another snippet here",
+                },
+            ]
+        }
+    }
+
+    mock_response = ComposioExecuteToolResponse(provider="COMPOSIO", data=mock_response_data)
+    base_node_setup._context.vellum_client.integrations.execute_integration_tool = MagicMock(return_value=mock_response)
+
+    # WHEN we run the node
+    outputs = base_node_setup.run()
+
+    # THEN text should include all results with snippet only when present
+    expected_text = """\
+Title: Result with snippet
+Snippet: This has a snippet
+URL: https://example1.com
+
+Title: Result without snippet
+URL: https://example2.com
+
+Title: Another with snippet
+Snippet: Another snippet here
+URL: https://example3.com\
+"""
+    assert outputs.text == expected_text
+
+    # AND URLs should include all links
+    assert outputs.urls == ["https://example1.com", "https://example2.com", "https://example3.com"]
+
+    assert len(outputs.results) == 3
+    assert outputs.results == [
+        {
+            "title": "Result with snippet",
+            "snippet": "This has a snippet",
+            "url": "https://example1.com",
+        },
+        {
+            "title": "Result without snippet",
+            "snippet": None,
+            "url": "https://example2.com",
+        },
+        {
+            "title": "Another with snippet",
+            "snippet": "Another snippet here",
+            "url": "https://example3.com",
+        },
+    ]

vellum/workflows/nodes/experimental/__init__.py

@@ -1,3 +1 @@
-from .tool_calling_node import ToolCallingNode
-
-__all__ = ["ToolCallingNode"]
+__all__: list[str] = []

vellum/workflows/nodes/mocks.py

@@ -1,14 +1,15 @@
 from functools import reduce
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Sequence, Type, Union
 
-from pydantic import ConfigDict, ValidationError
+from pydantic import ConfigDict, SerializationInfo, ValidationError, field_serializer, model_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.client.types.vellum_value import VellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.events.types import default_serializer
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
@@ -105,6 +106,25 @@ class MockNodeExecution(UniversalBaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    @model_serializer(mode="wrap")
+    def serialize_full_model(self, handler: Callable[[Any], Any], info: SerializationInfo) -> Dict[str, Any]:
+        """Serialize the model and add node_id field computed from then_outputs."""
+        serialized = handler(self)
+        serialized["node_id"] = str(self.then_outputs.__class__.__parent_class__.__id__)
+        return serialized
+
+    @field_serializer("then_outputs")
+    def serialize_then_outputs(self, then_outputs: BaseOutputs, _info: SerializationInfo) -> Dict[str, Any]:
+        return default_serializer(then_outputs)
+
+    @field_serializer("when_condition")
+    def serialize_when_condition(self, when_condition: BaseDescriptor, info: SerializationInfo) -> Dict[str, Any]:
+        # This allows `ee` to pass in workflow display context through the pydantic context
+        if isinstance(info.context, dict) and "serializer" in info.context and callable(info.context["serializer"]):
+            return info.context["serializer"](when_condition)
+
+        return default_serializer(when_condition)
+
     @staticmethod
     def validate_all(
         raw_mock_workflow_node_configs: Optional[List[Any]],
@@ -123,6 +143,7 @@ class MockNodeExecution(UniversalBaseModel):
             raise WorkflowInitializationException(
                 message="Failed to validate mock node executions",
                 code=WorkflowErrorCode.INVALID_INPUTS,
+                workflow_definition=workflow,
             ) from e
 
         nodes = {node.__id__: node for node in workflow.get_nodes()}
@@ -133,9 +154,9 @@ class MockNodeExecution(UniversalBaseModel):
         # We need to support the old way that the Vellum App's WorkflowRunner used to define Node Mocks in order to
         # avoid needing to update the mock resolution strategy that it and the frontend uses. The path towards
         # cleaning this up will go as follows:
-        # 1. Release Mock support in SDK-Enabled Workflows
-        # 2. Deprecate Mock support in non-SDK enabled Workflows, encouraging users to migrate to SDK-enabled Workflows
-        # 3. Remove the old mock resolution strategy
+        # 1. ✅ Release Mock support in SDK-Enabled Workflows
+        # 2. ✅ Deprecate Mock support in non-SDK enabled Workflows, encouraging users to migrate to SDK Workflows
+        # 3. ✅ Remove the old mock resolution strategy
         # 4. Update this SDK to handle the new mock resolution strategy with WorkflowValueDescriptors
         # 5. Cutover the Vellum App to the new mock resolution strategy
         # 6. Remove the old mock resolution strategy from this SDK
@@ -186,7 +207,10 @@ class MockNodeExecution(UniversalBaseModel):
             elif raw_logical_condition.operator == "!=":
                 return lhs.does_not_equal(rhs)
             else:
-                raise WorkflowInitializationException(f"Unsupported logical operator: {raw_logical_condition.operator}")
+                raise WorkflowInitializationException(
+                    message=f"Unsupported logical operator: {raw_logical_condition.operator}",
+                    workflow_definition=workflow,
+                )
 
         def _translate_raw_logical_expression_variable(
             raw_variable: _RawMockWorkflowNodeExecutionValuePointer,
@@ -197,7 +221,10 @@ class MockNodeExecution(UniversalBaseModel):
                 node = nodes[raw_variable.node_id]
                 return node.Execution.count
             else:
-                raise WorkflowInitializationException(f"Unsupported logical expression type: {raw_variable.type}")
+                raise WorkflowInitializationException(
+                    message=f"Unsupported logical expression type: {raw_variable.type}",
+                    workflow_definition=workflow,
+                )
 
         mock_node_executions = []
         for mock_workflow_node_config in mock_workflow_node_configs:
@@ -212,8 +239,10 @@ class MockNodeExecution(UniversalBaseModel):
                     for then_output in mock_execution.then_outputs:
                         node_output_name = node_output_name_by_id.get(then_output.output_id)
                         if node_output_name is None:
+                            node_id = mock_workflow_node_config.node_id
                             raise WorkflowInitializationException(
-                                f"Output {then_output.output_id} not found in node {mock_workflow_node_config.node_id}"
+                                message=f"Output {then_output.output_id} not found in node {node_id}",
+                                workflow_definition=workflow,
                             )
 
                         resolved_output_reference = _translate_raw_logical_expression_variable(then_output.value)
@@ -224,8 +253,10 @@ class MockNodeExecution(UniversalBaseModel):
                                 resolved_output_reference._value,
                             )
                         else:
+                            ref_type = type(resolved_output_reference)
                             raise WorkflowInitializationException(
-                                f"Unsupported resolved output reference type: {type(resolved_output_reference)}"
+                                message=f"Unsupported resolved output reference type: {ref_type}",
+                                workflow_definition=workflow,
                             )
 
                     mock_node_executions.append(

vellum/workflows/nodes/tests/test_mocks.py

@@ -54,7 +54,7 @@ def test_mocks__parse_from_app():
                     },
                     "then_outputs": [
                         {
-                            "output_id": "9e6dc5d3-8ea0-4346-8a2a-7cce5495755b",
+                            "output_id": "bbd955e8-b939-45f9-8617-f1405dddd2b3",
                             "value": {
                                 "id": "27006b2a-fa81-430c-a0b2-c66a9351fc68",
                                 "type": "CONSTANT_VALUE",
@@ -62,7 +62,7 @@ def test_mocks__parse_from_app():
                             },
                         },
                         {
-                            "output_id": "60305ffd-60b0-42aa-b54e-4fdae0f8c28a",
+                            "output_id": "e50bf6e9-417b-4c04-babb-94e7dccb9f2c",
                             "value": {
                                 "id": "4559c778-6e27-4cfe-a460-734ba62a5082",
                                 "type": "CONSTANT_VALUE",

vellum/workflows/nodes/tests/test_utils.py

@@ -1,5 +1,5 @@
 import pytest
-from typing import Any, List, Union
+from typing import Any, Dict, List, Union
 
 from pydantic import BaseModel
 
@@ -34,6 +34,9 @@ class Person(BaseModel):
             Json,
             {"name": "test", "args": [1, 2]},
         ),
+        ('{"name": "Alice", "age": 30}', dict, {"name": "Alice", "age": 30}),
+        ('{"name": "Alice", "age": 30}', Dict[str, Any], {"name": "Alice", "age": 30}),
+        ('{"name": "Alice", "age": 30}', Any, {"name": "Alice", "age": 30}),
         ("42", Union[int, str], 42),
         ("hello", Union[int, str], "hello"),
     ],
@@ -47,6 +50,9 @@ class Person(BaseModel):
         "list_of_str",
         "simple_json",
         "function_call_json",
+        "dict_type",
+        "dict_str_any_type",
+        "any_type",
         "union_int",
         "union_str",
     ],
@@ -86,23 +92,23 @@ def test_parse_type_from_str_list_of_models():
             List[str],
             NodeException,
             WorkflowErrorCode.INVALID_OUTPUTS,
-            "Invalid JSON Array format for result_as_str",
+            "Invalid JSON format: expected a valid JSON array",
         ),
         (
             "{invalid json}",
             Person,
             NodeException,
             WorkflowErrorCode.INVALID_OUTPUTS,
-            "Invalid JSON format for result_as_str",
+            "Invalid JSON format: unable to parse the provided data",
         ),
         (
             "{invalid json}",
             Json,
             NodeException,
             WorkflowErrorCode.INVALID_OUTPUTS,
-            "Invalid JSON format for result_as_str",
+            "Invalid JSON format: unable to parse the provided data",
         ),
-        ('{"name": "Alice"}', List[str], ValueError, None, "Expected a list of items for result_as_str, received dict"),
+        ('{"name": "Alice"}', List[str], ValueError, None, "Expected a list of items, but received dict"),
         ("data", object, ValueError, None, "Unsupported output type: <class 'object'>"),
     ],
     ids=[
@@ -150,3 +156,26 @@ def test_cast_to_output_type_none_value(output_type, expected_result):
     """Test that cast_to_output_type returns appropriate default values when None is provided."""
     result = cast_to_output_type(None, output_type)
     assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "input_value,expected_result",
+    [
+        ('{"name": "Alice", "age": 30}', {"name": "Alice", "age": 30}),
+        ("[1, 2, 3]", [1, 2, 3]),
+        ("invalid json", "invalid json"),
+        ([1, 2, 3], [1, 2, 3]),
+        ({"already": "dict"}, {"already": "dict"}),
+    ],
+    ids=[
+        "valid_json_object",
+        "valid_json_array",
+        "invalid_json_string",
+        "non_string_list",
+        "non_string_dict",
+    ],
+)
+def test_cast_to_output_type_any_json_parsing(input_value, expected_result):
+    """Test that cast_to_output_type attempts JSON parsing for Any type and falls back gracefully."""
+    result = cast_to_output_type(input_value, Any)
+    assert result == expected_result