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

vellum/workflows/nodes/displayable/bases/api_node/node.py

@@ -4,6 +4,8 @@ from requests import Request, RequestException, Session
 from requests.exceptions import JSONDecodeError
 
 from vellum.client import ApiError
+from vellum.client.core.request_options import RequestOptions
+from vellum.client.types.method_enum import MethodEnum
 from vellum.client.types.vellum_secret import VellumSecret as ClientVellumSecret
 from vellum.workflows.constants import APIRequestMethod
 from vellum.workflows.errors.types import WorkflowErrorCode
@@ -12,6 +14,7 @@ from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.types.core import Json, MergeBehavior, VellumSecret
 from vellum.workflows.types.generics import StateType
+from vellum.workflows.utils.hmac import sign_request_with_env_secret
 
 
 class BaseAPINode(BaseNode, Generic[StateType]):
@@ -19,20 +22,22 @@ class BaseAPINode(BaseNode, Generic[StateType]):
     Used to execute an API call.
 
     url: str - The URL to send the request to.
-    method: APIRequestMethod - The HTTP method to use for the request.
+    method: Union[APIRequestMethod, str] - The HTTP method to use for the request.
     data: Optional[str] - The data to send in the request body.
     json: Optional["JsonObject"] - The JSON data to send in the request body.
     headers: Optional[Dict[str, Union[str, VellumSecret]]] - The headers to send in the request.
+    timeout: Optional[int] - The timeout in seconds for the API request.
     """
 
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
     url: str = ""
-    method: Optional[APIRequestMethod] = APIRequestMethod.GET
+    method: Optional[Union[APIRequestMethod, MethodEnum]] = APIRequestMethod.GET
     data: Optional[str] = None
     json: Optional[Json] = None
     headers: Optional[Dict[str, Union[str, VellumSecret]]] = None
+    timeout: Optional[int] = None
 
     class Outputs(BaseOutputs):
         json: Optional[Json]
@@ -40,58 +45,107 @@ class BaseAPINode(BaseNode, Generic[StateType]):
         status_code: int
         text: str
 
+    def _normalize_http_method(self, method: Union[APIRequestMethod, MethodEnum]) -> str:
+        if isinstance(method, APIRequestMethod):
+            method_str = method.value
+        elif isinstance(method, str):
+            method_str = method.upper()
+            valid_methods = {m.value for m in APIRequestMethod}
+            if method_str not in valid_methods:
+                raise NodeException(f"Invalid HTTP method '{method}'", code=WorkflowErrorCode.INVALID_INPUTS)
+        else:
+            raise NodeException(
+                f"Method must be either APIRequestMethod enum or string, got {type(method)}",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+        return method_str
+
+    def _validate(self) -> None:
+        if not self.url or not isinstance(self.url, str) or not self.url.strip():
+            raise NodeException("URL is required and must be a non-empty string", code=WorkflowErrorCode.INVALID_INPUTS)
+
     def run(self) -> Outputs:
-        return self._run(method=self.method, url=self.url, data=self.data, json=self.json, headers=self.headers)
+        return self._run(
+            method=self.method, url=self.url, data=self.data, json=self.json, headers=self.headers, timeout=self.timeout
+        )
 
     def _run(
         self,
         url: str,
-        method: Optional[APIRequestMethod] = APIRequestMethod.GET,
+        method: Optional[Union[APIRequestMethod, MethodEnum]] = APIRequestMethod.GET,
         data: Optional[Union[str, Any]] = None,
         json: Any = None,
         headers: Any = None,
         bearer_token: Optional[VellumSecret] = None,
+        timeout: Optional[int] = None,
     ) -> Outputs:
+        self._validate()
+
+        normalized_method = self._normalize_http_method(method) if method is not None else APIRequestMethod.GET.value
+
         vellum_instance = False
         for header in headers or {}:
             if isinstance(headers[header], VellumSecret):
                 vellum_instance = True
         if vellum_instance or bearer_token:
-            return self._vellum_execute_api(bearer_token, json, headers, method, url)
+            return self._vellum_execute_api(bearer_token, json, headers, normalized_method, url, timeout)
         else:
-            return self._local_execute_api(data, headers, json, method, url)
+            return self._local_execute_api(data, headers, json, normalized_method, url, timeout)
 
-    def _local_execute_api(self, data, headers, json, method, url):
+    def _local_execute_api(self, data, headers, json, method, url, timeout):
+        headers = headers or {}
+        if "User-Agent" not in headers:
+            client_headers = self._context.vellum_client._client_wrapper.get_headers()
+            headers["User-Agent"] = client_headers.get("User-Agent")
         try:
             if data is not None:
-                prepped = Request(method=method.value, url=url, data=data, headers=headers).prepare()
+                prepped = Request(method=method, url=url, data=data, headers=headers).prepare()
             elif json is not None:
-                prepped = Request(method=method.value, url=url, json=json, headers=headers).prepare()
+                prepped = Request(method=method, url=url, json=json, headers=headers).prepare()
             else:
-                prepped = Request(method=method.value, url=url, headers=headers).prepare()
+                prepped = Request(method=method, url=url, headers=headers).prepare()
         except Exception as e:
             raise NodeException(f"Failed to prepare HTTP request: {e}", code=WorkflowErrorCode.PROVIDER_ERROR)
+
+        sign_request_with_env_secret(prepped)
+
         try:
             with Session() as session:
-                response = session.send(prepped)
+                response = session.send(prepped, timeout=timeout)
         except RequestException as e:
             raise NodeException(f"HTTP request failed: {e}", code=WorkflowErrorCode.PROVIDER_ERROR)
         try:
-            json = response.json()
+            json_response = response.json()
         except JSONDecodeError:
-            json = None
+            json_response = None
         return self.Outputs(
-            json=json,
+            json=json_response,
             headers={header: value for header, value in response.headers.items()},
             status_code=response.status_code,
             text=response.text,
         )
 
-    def _vellum_execute_api(self, bearer_token, data, headers, method, url):
+    def _vellum_execute_api(self, bearer_token, data, headers, method, url, timeout):
         client_vellum_secret = ClientVellumSecret(name=bearer_token.name) if bearer_token else None
+
+        headers = headers or {}
+        if "User-Agent" not in headers:
+            client_headers = self._context.vellum_client._client_wrapper.get_headers()
+            headers["User-Agent"] = client_headers.get("User-Agent")
+
+        # Create request_options if timeout is specified
+        request_options = None
+        if timeout is not None:
+            request_options = RequestOptions(timeout_in_seconds=timeout)
+
         try:
             vellum_response = self._context.vellum_client.execute_api(
-                url=url, method=method.value, body=data, headers=headers, bearer_token=client_vellum_secret
+                url=url,
+                method=method,
+                body=data,
+                headers=headers,
+                bearer_token=client_vellum_secret,
+                request_options=request_options,
             )
         except ApiError as e:
             raise NodeException(f"Failed to prepare HTTP request: {e}", code=WorkflowErrorCode.NODE_EXECUTION)

vellum/workflows/nodes/displayable/bases/api_node/tests/test_node.py

@@ -0,0 +1,150 @@
+import pytest
+import os
+from unittest.mock import patch
+
+from vellum.client.types.execute_api_response import ExecuteApiResponse
+from vellum.workflows.constants import APIRequestMethod
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.displayable.bases.api_node.node import BaseAPINode
+from vellum.workflows.types.core import VellumSecret
+
+
+@pytest.mark.parametrize("method_value", ["GET", "get", APIRequestMethod.GET])
+def test_api_node_with_string_method(method_value, vellum_client):
+    class TestAPINode(BaseAPINode):
+        method = method_value
+        url = "https://example.com"
+        headers = {"Authorization": VellumSecret(name="API_KEY")}
+
+    mock_response = ExecuteApiResponse(
+        json_={"status": "success"},
+        headers={"content-type": "application/json"},
+        status_code=200,
+        text='{"status": "success"}',
+    )
+    vellum_client.execute_api.return_value = mock_response
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert result.status_code == 200
+
+    vellum_client.execute_api.assert_called_once()
+    call_args = vellum_client.execute_api.call_args
+    assert call_args[1]["method"] == "GET"
+
+
+def test_api_node_with_invalid_method():
+    class TestAPINode(BaseAPINode):
+        method = "INVALID_METHOD"
+        url = "https://example.com"
+
+    node = TestAPINode()
+
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "Invalid HTTP method 'INVALID_METHOD'" == str(exc_info.value)
+
+
+def test_api_node_adds_user_agent_header_when_none_provided(requests_mock):
+    """
+    Tests that the API node adds User-Agent header when no headers are provided.
+    """
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.GET
+        url = "https://example.com/test"
+
+    response_mock = requests_mock.get(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert response_mock.last_request
+    assert "vellum-ai" in response_mock.last_request.headers.get("User-Agent", "")
+
+    assert result.status_code == 200
+
+
+def test_api_node_adds_user_agent_header_when_headers_provided_without_user_agent(requests_mock):
+    """
+    Tests that the API node adds User-Agent header when headers are provided but don't include User-Agent.
+    """
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.POST
+        url = "https://example.com/test"
+        headers = {"Content-Type": "application/json", "Custom-Header": "value"}
+        json = {"test": "data"}
+
+    response_mock = requests_mock.post(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert response_mock.last_request
+    assert "vellum-ai" in response_mock.last_request.headers.get("User-Agent", "")
+    assert response_mock.last_request.headers.get("Content-Type") == "application/json"
+    assert response_mock.last_request.headers.get("Custom-Header") == "value"
+
+    assert result.status_code == 200
+
+
+def test_api_node_preserves_custom_user_agent_header(requests_mock):
+    """
+    Tests that the API node preserves a custom User-Agent header if provided.
+    """
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.GET
+        url = "https://example.com/test"
+        headers = {"User-Agent": "Custom-Agent/1.0"}
+
+    response_mock = requests_mock.get(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    node = TestAPINode()
+    result = node.run()
+
+    assert response_mock.last_request
+    assert response_mock.last_request.headers.get("User-Agent") == "Custom-Agent/1.0"
+
+    assert result.status_code == 200
+
+
+def test_local_execute_api_with_hmac_secret(requests_mock):
+    """Test that _local_execute_api adds HMAC headers when VELLUM_HMAC_SECRET is set."""
+
+    class TestAPINode(BaseAPINode):
+        method = APIRequestMethod.POST
+        url = "https://example.com/test"
+        json = {"test": "data"}
+
+    response_mock = requests_mock.post(
+        "https://example.com/test",
+        json={"result": "success"},
+        status_code=200,
+    )
+
+    with patch.dict(os.environ, {"VELLUM_HMAC_SECRET": "test-secret"}):
+        node = TestAPINode()
+        result = node.run()
+
+    assert response_mock.last_request
+    assert "X-Vellum-Timestamp" in response_mock.last_request.headers
+    assert "X-Vellum-Signature" in response_mock.last_request.headers
+    assert result.status_code == 200

vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py

@@ -1,18 +1,33 @@
 from abc import abstractmethod
-from typing import ClassVar, Generator, Generic, Iterator, List, Optional, Union
+from itertools import chain
+from typing import ClassVar, Generator, Generic, Iterator, List, Optional, Union, cast
 
 from vellum import AdHocExecutePromptEvent, ExecutePromptEvent, PromptOutput
+from vellum.client.core import RequestOptions
 from vellum.client.core.api_error import ApiError
-from vellum.core import RequestOptions
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode, vellum_error_to_workflow_error
+from vellum.workflows.events.node import NodeExecutionStreamingEvent
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.generics import StateType
 
 
-class BasePromptNode(BaseNode, Generic[StateType]):
+def _contains_reference_to_output(reference: BaseDescriptor, target_reference: OutputReference) -> bool:
+    if reference == target_reference:
+        return True
+    if isinstance(reference, CoalesceExpression):
+        return _contains_reference_to_output(reference.lhs, target_reference) or _contains_reference_to_output(
+            reference.rhs, target_reference
+        )
+    return False
+
+
+class BasePromptNode(BaseNode[StateType], Generic[StateType]):
     # Inputs that are passed to the Prompt
     prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
 
@@ -46,15 +61,26 @@ class BasePromptNode(BaseNode, Generic[StateType]):
         except ApiError as e:
             self._handle_api_error(e)
 
-        # We don't use the INITIATED event anyway, so we can just skip it
-        # and use the exception handling to catch other api level errors
         try:
-            next(prompt_event_stream)
+            first_event = next(prompt_event_stream)
         except ApiError as e:
             self._handle_api_error(e)
+        else:
+            if first_event.state == "REJECTED":
+                workflow_error = vellum_error_to_workflow_error(first_event.error)
+                raise NodeException.of(workflow_error)
+            if first_event.state != "INITIATED":
+                prompt_event_stream = cast(
+                    Union[Iterator[AdHocExecutePromptEvent], Iterator[ExecutePromptEvent]],
+                    chain([first_event], prompt_event_stream),
+                )
 
         outputs: Optional[List[PromptOutput]] = None
+        exception: Optional[NodeException] = None
         for event in prompt_event_stream:
+            if exception:
+                continue
+
             if event.state == "INITIATED":
                 continue
             elif event.state == "STREAMING":
@@ -64,7 +90,10 @@ class BasePromptNode(BaseNode, Generic[StateType]):
                 yield BaseOutput(name="results", value=event.outputs)
             elif event.state == "REJECTED":
                 workflow_error = vellum_error_to_workflow_error(event.error)
-                raise NodeException.of(workflow_error)
+                exception = NodeException.of(workflow_error)
+
+        if exception:
+            raise exception
 
         return outputs
 
@@ -73,15 +102,47 @@ class BasePromptNode(BaseNode, Generic[StateType]):
             raise NodeException(
                 message=e.body.get("detail", "Provider credentials is missing or unavailable"),
                 code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+                raw_data=e.body,
             )
 
         elif e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
             raise NodeException(
                 message=e.body.get("detail", "Failed to execute Prompt"),
                 code=WorkflowErrorCode.INVALID_INPUTS,
+                raw_data=e.body,
             ) from e
 
         raise NodeException(
             message="Failed to execute Prompt",
             code=WorkflowErrorCode.INTERNAL_ERROR,
+            raw_data=e.body,
         ) from e
+
+    def __directly_emit_workflow_output__(
+        self,
+        event: NodeExecutionStreamingEvent,
+        workflow_output_descriptor: OutputReference,
+    ) -> bool:
+        if event.output.name != "results":
+            return False
+
+        if not isinstance(event.output.delta, str) and not event.output.is_initiated:
+            return False
+
+        target_nodes = [e.to_node for port in self.Ports for e in port.edges if e.to_node.__simulates_workflow_output__]
+        target_node_output = next(
+            (
+                o
+                for target_node in target_nodes
+                for o in target_node.Outputs
+                if o == workflow_output_descriptor.instance
+            ),
+            None,
+        )
+        if not target_node_output:
+            return False
+
+        if not isinstance(target_node_output.instance, BaseDescriptor):
+            return False
+
+        return _contains_reference_to_output(target_node_output.instance, event.node_definition.Outputs.text)