vellum-ai 1.2.5__py3-none-any.whl → 1.3.1__py3-none-any.whl

vellum/workflows/nodes/__init__.py

@@ -11,6 +11,7 @@ from vellum.workflows.nodes.displayable import (
     PromptDeploymentNode,
     SearchNode,
     SubworkflowDeploymentNode,
+    WebSearchNode,
 )
 from vellum.workflows.nodes.displayable.bases import (
     BaseInlinePromptNode as BaseInlinePromptNode,
@@ -43,4 +44,5 @@ __all__ = [
     "PromptDeploymentNode",
     "SearchNode",
     "SubworkflowDeploymentNode",
+    "WebSearchNode",
 ]

vellum/workflows/nodes/displayable/__init__.py

@@ -14,6 +14,7 @@ from .prompt_deployment_node import PromptDeploymentNode
 from .search_node import SearchNode
 from .subworkflow_deployment_node import SubworkflowDeploymentNode
 from .tool_calling_node import ToolCallingNode
+from .web_search_node import WebSearchNode
 
 __all__ = [
     "APINode",
@@ -31,5 +32,6 @@ __all__ = [
     "SearchNode",
     "TemplatingNode",
     "ToolCallingNode",
+    "WebSearchNode",
     "FinalOutputNode",
 ]

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

@@ -54,12 +54,23 @@ class FunctionCallNodeMixin:
                 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=DefaultStateEncoder)),
+                source=function_call_id,
             )
         )
         with state.__quiet__():

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

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

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

@@ -0,0 +1,133 @@
+import logging
+from typing import Any, ClassVar, Dict, List, Optional
+
+from requests import Request, RequestException, Session
+from requests.exceptions import JSONDecodeError
+
+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 WebSearchNode(BaseNode[StateType]):
+    """
+    Used to perform web search using SerpAPI.
+
+    query: str - The search query to execute
+    api_key: str - SerpAPI authentication key
+    num_results: int - Number of search results to return (default: 10)
+    location: Optional[str] - Geographic location filter for search
+    """
+
+    query: ClassVar[str] = ""
+    api_key: ClassVar[Optional[str]] = None
+    num_results: ClassVar[int] = 10
+    location: ClassVar[Optional[str]] = None
+
+    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(self) -> None:
+        """Validate node inputs."""
+        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
+            )
+
+        if self.api_key is None:
+            raise NodeException("API key is required", code=WorkflowErrorCode.INVALID_INPUTS)
+
+        if not isinstance(self.num_results, int) or self.num_results <= 0:
+            raise NodeException("num_results must be a positive integer", code=WorkflowErrorCode.INVALID_INPUTS)
+
+    def run(self) -> Outputs:
+        """Run the WebSearchNode to perform web search via SerpAPI."""
+        self._validate()
+
+        api_key_value = self.api_key
+
+        params = {
+            "q": self.query,
+            "api_key": api_key_value,
+            "num": self.num_results,
+            "engine": "google",
+        }
+
+        if self.location:
+            params["location"] = self.location
+
+        headers = {}
+        client_headers = self._context.vellum_client._client_wrapper.get_headers()
+        headers["User-Agent"] = client_headers.get("User-Agent")
+
+        try:
+            prepped = Request(method="GET", url="https://serpapi.com/search", params=params, headers=headers).prepare()
+        except Exception as e:
+            logger.exception("Failed to prepare SerpAPI request")
+            raise NodeException(f"Failed to prepare HTTP request: {e}", code=WorkflowErrorCode.PROVIDER_ERROR) from e
+
+        try:
+            with Session() as session:
+                response = session.send(prepped, timeout=30)
+        except RequestException as e:
+            logger.exception("SerpAPI request failed")
+            raise NodeException(f"HTTP request failed: {e}", code=WorkflowErrorCode.PROVIDER_ERROR) from e
+
+        if response.status_code == 401:
+            logger.error("SerpAPI authentication failed")
+            raise NodeException("Invalid API key", code=WorkflowErrorCode.INVALID_INPUTS)
+        elif response.status_code == 429:
+            logger.warning("SerpAPI rate limit exceeded")
+            raise NodeException("Rate limit exceeded", code=WorkflowErrorCode.PROVIDER_ERROR)
+        elif response.status_code >= 400:
+            logger.error(f"SerpAPI returned error status: {response.status_code}")
+            raise NodeException(f"SerpAPI error: HTTP {response.status_code}", code=WorkflowErrorCode.PROVIDER_ERROR)
+
+        try:
+            json_response = response.json()
+        except JSONDecodeError as e:
+            logger.exception("Failed to parse SerpAPI response as JSON")
+            raise NodeException(
+                f"Invalid JSON response from SerpAPI: {e}", code=WorkflowErrorCode.PROVIDER_ERROR
+            ) from e
+
+        if "error" in json_response:
+            error_msg = json_response["error"]
+            logger.error(f"SerpAPI returned error: {error_msg}")
+            raise NodeException(f"SerpAPI error: {error_msg}", code=WorkflowErrorCode.PROVIDER_ERROR)
+
+        organic_results = json_response.get("organic_results", [])
+
+        text_results = []
+        urls = []
+
+        for result in organic_results:
+            title = result.get("title", "")
+            snippet = result.get("snippet", "")
+            link = result.get("link", "")
+
+            if title and snippet:
+                text_results.append(f"{title}: {snippet}")
+            elif title:
+                text_results.append(title)
+            elif snippet:
+                text_results.append(snippet)
+
+            if link:
+                urls.append(link)
+
+        return self.Outputs(text="\n\n".join(text_results), urls=urls, results=organic_results)

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

@@ -0,0 +1,319 @@
+import pytest
+from unittest.mock import MagicMock
+
+import requests
+
+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
+        api_key: str
+        num_results: int
+
+    class State(BaseState):
+        pass
+
+    class TestableWebSearchNode(WebSearchNode):
+        query = Inputs.query
+        api_key = Inputs.api_key
+        num_results = Inputs.num_results
+
+    state = State(meta=StateMeta(workflow_inputs=Inputs(query="test query", api_key="test_api_key", num_results=3)))
+    context = MagicMock()
+    context.vellum_client = vellum_client
+    node = TestableWebSearchNode(state=state, context=context)
+    return node
+
+
+def test_successful_search_with_results(base_node_setup, requests_mock):
+    """Test successful SerpAPI search with typical organic results."""
+    # GIVEN a mock SerpAPI response with organic results
+    mock_response = {
+        "organic_results": [
+            {
+                "title": "First Result",
+                "snippet": "This is the first search result snippet",
+                "link": "https://example1.com",
+                "position": 1,
+            },
+            {
+                "title": "Second Result",
+                "snippet": "This is the second search result snippet",
+                "link": "https://example2.com",
+                "position": 2,
+            },
+            {
+                "title": "Third Result",
+                "snippet": "This is the third search result snippet",
+                "link": "https://example3.com",
+                "position": 3,
+            },
+        ]
+    }
+
+    requests_mock.get("https://serpapi.com/search", json=mock_response)
+
+    # WHEN we run the node
+    outputs = base_node_setup.run()
+
+    # THEN the text output should be properly formatted
+    expected_text = (
+        "First Result: This is the first search result snippet\n\n"
+        "Second Result: This is the second search result snippet\n\n"
+        "Third Result: This is the third search result snippet"
+    )
+    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 preserved
+    assert outputs.results == mock_response["organic_results"]
+
+    # AND the request should have the correct parameters
+    assert requests_mock.last_request.qs == {
+        "q": ["test query"],
+        "api_key": ["test_api_key"],
+        "num": ["3"],
+        "engine": ["google"],
+    }
+
+
+def test_search_with_location_parameter(base_node_setup, requests_mock):
+    """Test that location parameter is properly passed to SerpAPI."""
+    # GIVEN a location parameter is set
+    base_node_setup.location = "New York, NY"
+
+    requests_mock.get("https://serpapi.com/search", json={"organic_results": []})
+
+    # WHEN we run the node
+    base_node_setup.run()
+
+    # THEN the location parameter should be included (URL encoding may lowercase)
+    assert "location" in requests_mock.last_request.qs
+    assert requests_mock.last_request.qs["location"][0].lower() == "new york, ny"
+
+
+def test_authentication_error_401(base_node_setup, requests_mock):
+    """Test 401 authentication error raises NodeException with INVALID_INPUTS."""
+    # GIVEN SerpAPI returns a 401 authentication error
+    requests_mock.get("https://serpapi.com/search", status_code=401)
+
+    # 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.INVALID_INPUTS
+    assert "Invalid API key" in str(exc_info.value)
+
+
+def test_rate_limit_error_429(base_node_setup, requests_mock):
+    """Test 429 rate limit error raises NodeException with PROVIDER_ERROR."""
+    # GIVEN SerpAPI returns a 429 rate limit error
+    requests_mock.get("https://serpapi.com/search", status_code=429)
+
+    # 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 "Rate limit exceeded" in str(exc_info.value)
+
+
+def test_server_error_500(base_node_setup, requests_mock):
+    """Test 500+ server errors raise NodeException with PROVIDER_ERROR."""
+    # GIVEN SerpAPI returns a 500 server error
+    requests_mock.get("https://serpapi.com/search", status_code=500)
+
+    # 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 "SerpAPI error: HTTP 500" in str(exc_info.value)
+
+
+def test_invalid_json_response(base_node_setup, requests_mock):
+    """Test non-JSON response raises appropriate NodeException."""
+    # GIVEN SerpAPI returns non-JSON content
+    requests_mock.get("https://serpapi.com/search", text="Not JSON")
+
+    # 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 "Invalid JSON response" in str(exc_info.value)
+
+
+def test_serpapi_error_in_response(base_node_setup, requests_mock):
+    """Test SerpAPI error field in response raises NodeException."""
+    # GIVEN SerpAPI returns an error in the response body
+    requests_mock.get("https://serpapi.com/search", json={"error": "Invalid search parameters"})
+
+    # 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 "Invalid search parameters" in str(exc_info.value)
+
+
+def test_empty_query_validation(vellum_client):
+    """Test empty query raises validation error."""
+
+    # GIVEN a node with an empty query
+    class TestNode(WebSearchNode):
+        query = ""
+        api_key = "test_key"
+        num_results = 10
+
+    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 "Query is required" in str(exc_info.value)
+
+
+def test_missing_api_key_validation(vellum_client):
+    """Test missing API key raises validation error."""
+
+    # GIVEN a node with no API key
+    class TestNode(WebSearchNode):
+        query = "test query"
+        api_key = None
+        num_results = 10
+
+    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 "API key is required" in str(exc_info.value)
+
+
+def test_invalid_num_results_validation(vellum_client):
+    """Test invalid num_results raises validation error."""
+
+    # GIVEN a node with invalid num_results
+    class TestNode(WebSearchNode):
+        query = "test query"
+        api_key = "test_key"
+        num_results = -1
+
+    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 "num_results must be a positive integer" in str(exc_info.value)
+
+
+def test_empty_organic_results(base_node_setup, requests_mock):
+    """Test handling of empty search results."""
+    # GIVEN SerpAPI returns no organic results
+    requests_mock.get("https://serpapi.com/search", json={"organic_results": []})
+
+    # 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, requests_mock):
+    """Test handling of missing title, snippet, or link fields."""
+    # GIVEN SerpAPI returns results with missing fields
+    mock_response = {
+        "organic_results": [
+            {
+                "title": "Only Title",
+                # Missing snippet and link
+            },
+            {
+                "snippet": "Only snippet, no title or link"
+                # Missing title and link
+            },
+            {
+                "title": "Title with link",
+                "link": "https://example.com",
+                # Missing snippet
+            },
+            {
+                # All fields missing - should be skipped
+                "position": 4
+            },
+        ]
+    }
+
+    requests_mock.get("https://serpapi.com/search", json=mock_response)
+
+    # WHEN we run the node
+    outputs = base_node_setup.run()
+
+    # THEN text should handle missing fields gracefully
+    expected_text = "Only Title\n\n" "Only snippet, no title or link\n\n" "Title with link"
+    assert outputs.text == expected_text
+
+    # AND URLs should only include valid links
+    assert outputs.urls == ["https://example.com"]
+
+
+def test_request_timeout_handling(base_node_setup, requests_mock):
+    """Test network timeout raises appropriate error."""
+    # GIVEN a network timeout occurs
+    requests_mock.get("https://serpapi.com/search", exc=requests.exceptions.Timeout("Connection timed out"))
+
+    # WHEN we run the node
+    with pytest.raises(NodeException) as exc_info:
+        base_node_setup.run()
+
+    # THEN it should raise a provider error
+    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_ERROR
+    assert "HTTP request failed" in str(exc_info.value)
+
+
+def test_user_agent_header_included(base_node_setup, requests_mock):
+    """Test that User-Agent header from vellum_client is included."""
+    # GIVEN a successful request
+    requests_mock.get("https://serpapi.com/search", json={"organic_results": []})
+
+    # WHEN we run the node
+    base_node_setup.run()
+
+    # THEN the User-Agent header should be included
+    assert requests_mock.last_request.headers["User-Agent"] == "vellum-python-sdk/1.0.0"

vellum/workflows/resolvers/base.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
 from uuid import UUID
-from typing import TYPE_CHECKING, Iterator, Optional, Type
+from typing import TYPE_CHECKING, Iterator, Optional, Type, Union
 
 from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.resolvers.types import LoadStateResult
 from vellum.workflows.state.base import BaseState
 
 if TYPE_CHECKING:
@@ -28,5 +29,5 @@ class BaseWorkflowResolver(ABC):
         pass
 
     @abstractmethod
-    def load_state(self, previous_execution_id: Optional[UUID] = None) -> Optional[BaseState]:
+    def load_state(self, previous_execution_id: Optional[Union[UUID, str]] = None) -> Optional[LoadStateResult]:
         pass

vellum/workflows/resolvers/resolver.py

@@ -1,10 +1,13 @@
 import logging
 from uuid import UUID
-from typing import Iterator, Optional
+from typing import Iterator, List, Optional, Tuple, Union
 
+from vellum.client.types.vellum_span import VellumSpan
+from vellum.client.types.workflow_execution_initiated_event import WorkflowExecutionInitiatedEvent
 from vellum.workflows.events.workflow import WorkflowEvent
 from vellum.workflows.resolvers.base import BaseWorkflowResolver
-from vellum.workflows.state.base import BaseState, StateMeta
+from vellum.workflows.resolvers.types import LoadStateResult
+from vellum.workflows.state.base import BaseState
 
 logger = logging.getLogger(__name__)
 
@@ -16,7 +19,42 @@ class VellumResolver(BaseWorkflowResolver):
     def get_state_snapshot_history(self) -> Iterator[BaseState]:
         return iter([])
 
-    def load_state(self, previous_execution_id: Optional[UUID] = None) -> Optional[BaseState]:
+    def _find_previous_and_root_span(
+        self, execution_id: str, spans: List[VellumSpan]
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+        previous_trace_id: Optional[str] = None
+        root_trace_id: Optional[str] = None
+        previous_span_id: Optional[str] = None
+        root_span_id: Optional[str] = None
+
+        for span in spans:
+            # Look for workflow execution spans with matching ID first
+            if span.name == "workflow.execution" and span.span_id == execution_id:
+                # Find the WorkflowExecutionInitiatedEvent in the span's events
+                initiated_event = next(
+                    (event for event in span.events if isinstance(event, WorkflowExecutionInitiatedEvent)), None
+                )
+                if initiated_event:
+                    previous_trace_id = initiated_event.trace_id
+                    previous_span_id = initiated_event.span_id
+                    links = initiated_event.links
+                    if links:
+                        root_span = next((link for link in links if link.type == "ROOT_SPAN"), None)
+                        if root_span:
+                            root_trace_id = root_span.trace_id
+                            root_span_id = root_span.span_context.span_id
+                    else:
+                        # no links means this is the first execution
+                        root_trace_id = initiated_event.trace_id
+                        root_span_id = initiated_event.span_id
+                    break
+
+        return previous_trace_id, root_trace_id, previous_span_id, root_span_id
+
+    def load_state(self, previous_execution_id: Optional[Union[UUID, str]] = None) -> Optional[LoadStateResult]:
+        if isinstance(previous_execution_id, UUID):
+            previous_execution_id = str(previous_execution_id)
+
         if previous_execution_id is None:
             return None
 
@@ -26,17 +64,34 @@ class VellumResolver(BaseWorkflowResolver):
 
         client = self._context.vellum_client
         response = client.workflow_executions.retrieve_workflow_execution_detail(
-            execution_id=str(previous_execution_id),
+            execution_id=previous_execution_id,
         )
 
         if response.state is None:
             return None
 
-        meta = StateMeta.model_validate(response.state.pop("meta"))
+        previous_trace_id, root_trace_id, previous_span_id, root_span_id = self._find_previous_and_root_span(
+            previous_execution_id, response.spans
+        )
+
+        if previous_trace_id is None or root_trace_id is None or previous_span_id is None or root_span_id is None:
+            logger.warning("Could not find required execution events for state loading")
+            return None
+
+        if "meta" in response.state:
+            response.state.pop("meta")
 
         if self._workflow_class:
             state_class = self._workflow_class.get_state_class()
-            return state_class(**response.state, meta=meta)
+            state = state_class(**response.state)
         else:
             logger.warning("No workflow class registered, falling back to BaseState")
-            return BaseState(**response.state, meta=meta)
+            state = BaseState(**response.state)
+
+        return LoadStateResult(
+            state=state,
+            previous_trace_id=previous_trace_id,
+            previous_span_id=previous_span_id,
+            root_trace_id=root_trace_id,
+            root_span_id=root_span_id,
+        )