vellum-ai 1.2.4__py3-none-any.whl → 1.3.0__py3-none-any.whl

vellum/types/named_test_case_audio_variable_value_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_audio_variable_value_request import *

vellum/types/named_test_case_document_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_document_variable_value import *

vellum/types/named_test_case_document_variable_value_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_document_variable_value_request import *

vellum/types/named_test_case_image_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_image_variable_value import *

vellum/types/named_test_case_image_variable_value_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_image_variable_value_request import *

vellum/types/named_test_case_video_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_video_variable_value import *

vellum/types/named_test_case_video_variable_value_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.named_test_case_video_variable_value_request import *

vellum/types/scenario_input_audio_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.scenario_input_audio_variable_value import *

vellum/types/scenario_input_document_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.scenario_input_document_variable_value import *

vellum/types/scenario_input_image_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.scenario_input_image_variable_value import *

vellum/types/scenario_input_video_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.scenario_input_video_variable_value import *

vellum/types/test_case_audio_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.test_case_audio_variable_value import *

vellum/types/test_case_document_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.test_case_document_variable_value import *

vellum/types/test_case_image_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.test_case_image_variable_value import *

vellum/types/test_case_video_variable_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.test_case_video_variable_value import *

vellum/types/video_input.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.video_input import *

vellum/types/workflow_request_audio_input_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.workflow_request_audio_input_request import *

vellum/types/workflow_request_document_input_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.workflow_request_document_input_request import *

vellum/types/workflow_request_image_input_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.workflow_request_image_input_request import *

vellum/types/workflow_request_video_input_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.workflow_request_video_input_request import *

vellum/workflows/events/types.py

@@ -7,7 +7,6 @@ from pydantic import Field, GetCoreSchemaHandler, Tag, ValidationInfo
 from pydantic_core import CoreSchema, core_schema
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
-from vellum.client.types.span_link import SpanLink
 from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.definition import VellumCodeResourceDefinition
 from vellum.workflows.types.utils import datetime_now
@@ -86,6 +85,12 @@ class ExternalParentContext(BaseParentContext):
     type: Literal["EXTERNAL"] = "EXTERNAL"
 
 
+class SpanLink(UniversalBaseModel):
+    trace_id: str
+    type: Literal["TRIGGERED_BY", "PREVIOUS_SPAN", "ROOT_SPAN"]
+    span_context: "ParentContext"
+
+
 def _cast_parent_context_discriminator(v: Any) -> Any:
     if v in PARENT_CONTEXT_TYPES:
         return v

vellum/workflows/integrations/tests/test_mcp_service.py

@@ -1,10 +1,11 @@
+import pytest
 import asyncio
 import json
 from unittest import mock
 
 from vellum.workflows.constants import AuthorizationType
 from vellum.workflows.integrations.mcp_service import MCPHttpClient, MCPService
-from vellum.workflows.types.definition import MCPServer
+from vellum.workflows.types.definition import MCPServer, MCPToolDefinition
 
 
 def test_mcp_http_client_sse_response():
@@ -118,3 +119,107 @@ def test_mcp_service_api_key_auth():
 
     # THEN the custom API key header should be set correctly
     assert headers == {"X-API-Key": "api-key-123"}
+
+
+@pytest.mark.asyncio
+async def test_mcp_http_client_empty_response():
+    """Test that empty responses are handled gracefully"""
+    # GIVEN a mock response that returns empty content
+    mock_response = mock.Mock()
+    mock_response.headers = {"content-type": "application/json"}
+    mock_response.text = ""
+
+    # AND a mock httpx client that returns this response
+    with mock.patch("vellum.workflows.integrations.mcp_service.httpx.AsyncClient") as mock_client_class:
+        mock_client = mock.AsyncMock()
+        mock_client.post.return_value = mock_response
+        mock_client_class.return_value = mock_client
+
+        # WHEN we call initialize with an empty response
+        # THEN it should raise an exception about empty response
+        async with MCPHttpClient("https://test.server.com", {}) as client:
+            with pytest.raises(Exception, match="Empty response received from server"):
+                await client.initialize()
+
+
+@pytest.mark.asyncio
+async def test_mcp_http_client_invalid_sse_json():
+    """Test that invalid JSON in SSE data is handled"""
+    # GIVEN an SSE response with invalid JSON
+    invalid_sse = """event: message
+data: {invalid json}
+
+"""
+
+    mock_response = mock.Mock()
+    mock_response.headers = {"content-type": "text/event-stream"}
+    mock_response.text = invalid_sse
+
+    with mock.patch("vellum.workflows.integrations.mcp_service.httpx.AsyncClient") as mock_client_class:
+        mock_client = mock.AsyncMock()
+        mock_client.post.return_value = mock_response
+        mock_client_class.return_value = mock_client
+
+        # WHEN we call initialize with invalid SSE data
+        # THEN it should raise an exception about no valid JSON
+        async with MCPHttpClient("https://test.server.com", {}) as client:
+            with pytest.raises(Exception, match="No valid JSON data found in SSE response"):
+                await client.initialize()
+
+
+def test_mcp_service_hydrate_tool_definitions():
+    """Test tool definition hydration with SSE responses"""
+    # GIVEN an MCP server configuration
+    sample_mcp_server = MCPServer(
+        name="test-server",
+        url="https://test.mcp.server.com/mcp",
+        authorization_type=AuthorizationType.BEARER_TOKEN,
+        bearer_token_value="test-token-123",
+    )
+
+    # AND a mock MCP service that returns tools via SSE
+    with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run:
+        mock_run.return_value = [
+            {
+                "name": "resolve-library-id",
+                "description": "Resolves library names to IDs",
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {"libraryName": {"type": "string"}},
+                    "required": ["libraryName"],
+                },
+            }
+        ]
+
+        # WHEN we hydrate tool definitions
+        service = MCPService()
+        tool_definitions = service.hydrate_tool_definitions(sample_mcp_server)
+
+        # THEN we should get properly formatted MCPToolDefinition objects
+        assert len(tool_definitions) == 1
+        assert isinstance(tool_definitions[0], MCPToolDefinition)
+        assert tool_definitions[0].name == "resolve-library-id"
+        assert tool_definitions[0].description == "Resolves library names to IDs"
+        assert tool_definitions[0].server == sample_mcp_server
+        assert tool_definitions[0].parameters == {
+            "type": "object",
+            "properties": {"libraryName": {"type": "string"}},
+            "required": ["libraryName"],
+        }
+
+
+def test_mcp_service_list_tools_handles_errors():
+    """Test that SSE parsing errors are handled gracefully"""
+    # GIVEN an MCP server configuration
+    sample_mcp_server = MCPServer(name="test-server", url="https://test.mcp.server.com/mcp")
+
+    # AND a mock that raises an exception during SSE parsing
+    with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run:
+        mock_run.side_effect = Exception("SSE parsing failed")
+
+        # WHEN we try to list tools
+        service = MCPService()
+        tools = service.list_tools(sample_mcp_server)
+
+        # THEN we should get an empty list instead of crashing
+        assert tools == []

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/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/resolvers/base.py

@@ -1,11 +1,25 @@
 from abc import ABC, abstractmethod
-from typing import Iterator
+from uuid import UUID
+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:
+    from vellum.workflows.state.context import WorkflowContext
+    from vellum.workflows.workflows.base import BaseWorkflow
+
 
 class BaseWorkflowResolver(ABC):
+    def __init__(self):
+        self._context: Optional["WorkflowContext"] = None
+        self._workflow_class: Optional[Type["BaseWorkflow"]] = None
+
+    def register_workflow_instance(self, workflow_instance: "BaseWorkflow") -> None:
+        self._workflow_class = type(workflow_instance)
+        self._context = workflow_instance.context
+
     @abstractmethod
     def get_latest_execution_events(self) -> Iterator[WorkflowEvent]:
         pass
@@ -13,3 +27,7 @@ class BaseWorkflowResolver(ABC):
     @abstractmethod
     def get_state_snapshot_history(self) -> Iterator[BaseState]:
         pass
+
+    @abstractmethod
+    def load_state(self, previous_execution_id: Optional[Union[UUID, str]] = None) -> Optional[LoadStateResult]:
+        pass

vellum/workflows/resolvers/resolver.py

@@ -0,0 +1,97 @@
+import logging
+from uuid import UUID
+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.resolvers.types import LoadStateResult
+from vellum.workflows.state.base import BaseState
+
+logger = logging.getLogger(__name__)
+
+
+class VellumResolver(BaseWorkflowResolver):
+    def get_latest_execution_events(self) -> Iterator[WorkflowEvent]:
+        return iter([])
+
+    def get_state_snapshot_history(self) -> Iterator[BaseState]:
+        return iter([])
+
+    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
+
+        if not self._context:
+            logger.warning("Cannot load state: No workflow context registered")
+            return None
+
+        client = self._context.vellum_client
+        response = client.workflow_executions.retrieve_workflow_execution_detail(
+            execution_id=previous_execution_id,
+        )
+
+        if response.state is None:
+            return None
+
+        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()
+            state = state_class(**response.state)
+        else:
+            logger.warning("No workflow class registered, falling back to BaseState")
+            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,
+        )