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

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

@@ -1,11 +1,12 @@
-from typing import Optional, Union
+from typing import Generic, Optional, Union
 
 from vellum.workflows.constants import AuthorizationType
 from vellum.workflows.nodes.displayable.bases.api_node import BaseAPINode
 from vellum.workflows.types.core import MergeBehavior, VellumSecret
+from vellum.workflows.types.generics import StateType
 
 
-class APINode(BaseAPINode):
+class APINode(BaseAPINode[StateType], Generic[StateType]):
     """
     Used to execute an API call. This node exists to be backwards compatible with Vellum's API Node, and for most cases,
     you should extend from `BaseAPINode` directly.
@@ -22,6 +23,10 @@ class APINode(BaseAPINode):
     for the bearer token authorization.
     """
 
+    class Display(BaseAPINode.Display):
+        icon = "vellum:icon:signal-stream"
+        color = "lightBlue"
+
     authorization_type: Optional[AuthorizationType] = None
     api_key_header_key: Optional[str] = None
     api_key_header_value: Optional[Union[str, VellumSecret]] = None
@@ -31,6 +36,8 @@ class APINode(BaseAPINode):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
     def run(self) -> BaseAPINode.Outputs:
+        self._validate()
+
         headers = self.headers or {}
         header_overrides = {}
         bearer_token = None
@@ -51,7 +58,9 @@ class APINode(BaseAPINode):
         final_headers = {**headers, **header_overrides}
 
         vellum_client_wrapper = self._context.vellum_client._client_wrapper
-        if self.url.startswith(vellum_client_wrapper._environment.default) and "X-API-Key" not in final_headers:
+        if self.url.startswith(vellum_client_wrapper._environment.default) and (
+            "X-API-Key" not in final_headers and "X_API_KEY" not in final_headers
+        ):
             final_headers["X-API-Key"] = vellum_client_wrapper.api_key
 
         return self._run(
@@ -61,4 +70,5 @@ class APINode(BaseAPINode):
             json=self.json,
             headers=final_headers,
             bearer_token=bearer_token,
+            timeout=self.timeout,
         )

vellum/workflows/nodes/displayable/api_node/tests/test_api_node.py

@@ -3,6 +3,7 @@ import pytest
 from vellum import ExecuteApiResponse, VellumSecret as ClientVellumSecret
 from vellum.client.core.api_error import ApiError
 from vellum.workflows.constants import APIRequestMethod, AuthorizationType
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes import APINode
 from vellum.workflows.types.core import VellumSecret
@@ -15,6 +16,7 @@ def test_run_workflow__secrets(vellum_client):
         json_={"data": [1, 2, 3]},
         headers={"X-Response-Header": "bar"},
     )
+    vellum_client._client_wrapper.get_headers.return_value = {"User-Agent": "vellum-ai/1.0.6"}
 
     class SimpleBaseAPINode(APINode):
         method = APIRequestMethod.POST
@@ -34,7 +36,9 @@ def test_run_workflow__secrets(vellum_client):
     assert vellum_client.execute_api.call_count == 1
     assert vellum_client.execute_api.call_args.kwargs["url"] == "https://example.vellum.ai"
     assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
-    assert vellum_client.execute_api.call_args.kwargs["headers"] == {"X-Test-Header": "foo"}
+    headers = vellum_client.execute_api.call_args.kwargs["headers"]
+    assert headers["X-Test-Header"] == "foo"
+    assert "vellum-ai" in headers.get("User-Agent", "")
     bearer_token = vellum_client.execute_api.call_args.kwargs["bearer_token"]
     assert bearer_token == ClientVellumSecret(name="secret")
     assert terminal.headers == {"X-Response-Header": "bar"}
@@ -43,6 +47,7 @@ def test_run_workflow__secrets(vellum_client):
 def test_api_node_raises_error_when_api_call_fails(vellum_client):
     # GIVEN an API call that fails
     vellum_client.execute_api.side_effect = ApiError(status_code=400, body="API Error")
+    vellum_client._client_wrapper.get_headers.return_value = {"User-Agent": "vellum-ai/1.0.6"}
 
     class SimpleAPINode(APINode):
         method = APIRequestMethod.GET
@@ -69,7 +74,9 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
     assert vellum_client.execute_api.call_count == 1
     assert vellum_client.execute_api.call_args.kwargs["url"] == "https://example.vellum.ai"
     assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
-    assert vellum_client.execute_api.call_args.kwargs["headers"] == {"X-Test-Header": "foo"}
+    headers = vellum_client.execute_api.call_args.kwargs["headers"]
+    assert headers["X-Test-Header"] == "foo"
+    assert "vellum-ai" in headers.get("User-Agent", "")
 
 
 def test_api_node_defaults_to_get_method(vellum_client):
@@ -158,3 +165,145 @@ def test_api_node__detects_client_environment_urls__does_not_override_headers(
     # AND the vellum API should have been called with the correct headers
     assert mock_response.last_request
     assert mock_response.last_request.headers["X-API-Key"] == "vellum-api-key-5678"
+
+
+def test_api_node__detects_client_environment_urls__legacy_does_not_override_headers(
+    mock_httpx_transport, mock_requests, monkeypatch
+):
+    # GIVEN an API node with a URL pointing back to Vellum
+    class SimpleAPINodeToVellum(APINode):
+        url = "https://api.vellum.ai"
+        headers = {
+            "X_API_KEY": "vellum-api-key-5678",
+        }
+
+    # AND a mock request sent to the Vellum API would return a 200
+    mock_response = mock_requests.get(
+        "https://api.vellum.ai",
+        status_code=200,
+        json={"data": [1, 2, 3]},
+    )
+
+    # AND an api key is set
+    monkeypatch.setenv("VELLUM_API_KEY", "vellum-api-key-1234")
+
+    # WHEN we run the node
+    node = SimpleAPINodeToVellum()
+    node.run()
+
+    # THEN the execute_api method should not have been called
+    mock_httpx_transport.handle_request.assert_not_called()
+
+    # AND the vellum API should have been called with the correct headers
+    assert mock_response.last_request
+    assert mock_response.last_request.headers["X_API_KEY"] == "vellum-api-key-5678"
+
+
+def test_api_node_raises_error_for_empty_url():
+    class APINodeWithEmptyURL(APINode):
+        method = APIRequestMethod.GET
+        url = ""
+
+    node = APINodeWithEmptyURL()
+
+    with pytest.raises(NodeException) as excinfo:
+        node.run()
+
+    assert excinfo.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "URL is required and must be a non-empty string" in str(excinfo.value)
+
+
+def test_api_node_raises_error_for_none_url():
+    class APINodeWithNoneURL(APINode):
+        method = APIRequestMethod.GET
+        url = None  # type: ignore
+
+    node = APINodeWithNoneURL()
+
+    with pytest.raises(NodeException) as excinfo:
+        node.run()
+
+    assert excinfo.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "URL is required and must be a non-empty string" in str(excinfo.value)
+
+
+def test_api_node_raises_error_for_whitespace_url():
+    class APINodeWithWhitespaceURL(APINode):
+        method = APIRequestMethod.GET
+        url = "   "
+
+    node = APINodeWithWhitespaceURL()
+
+    with pytest.raises(NodeException) as excinfo:
+        node.run()
+
+    assert excinfo.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "URL is required and must be a non-empty string" in str(excinfo.value)
+
+
+def test_api_node_passes_timeout_to_vellum_client(vellum_client):
+    """Test that timeout is passed to vellum client as RequestOptions."""
+    vellum_client.execute_api.return_value = ExecuteApiResponse(
+        status_code=200,
+        text='{"result": "success"}',
+        json_={"result": "success"},
+        headers={"content-type": "application/json"},
+    )
+
+    # GIVEN an API node configured with a timeout value of 25 seconds
+    class APINodeWithTimeout(APINode):
+        method = APIRequestMethod.GET
+        authorization_type = AuthorizationType.BEARER_TOKEN
+        url = "https://example.com"
+        timeout = 25
+        bearer_token_value = VellumSecret(name="secret")
+
+    # WHEN the API node is executed
+    node = APINodeWithTimeout()
+    node.run()
+
+    # THEN the vellum client should be called exactly once
+    assert vellum_client.execute_api.call_count == 1
+    call_args = vellum_client.execute_api.call_args
+    request_options = call_args.kwargs["request_options"]
+    # AND the call should include RequestOptions with the correct timeout
+    assert request_options is not None
+    assert request_options["timeout_in_seconds"] == 25
+
+
+def test_api_node_supports_state_type_parameter(vellum_client):
+    """Test that APINode supports StateType generic parameter."""
+
+    # GIVEN a state class
+    from vellum.workflows.state.base import BaseState
+
+    class TestState(BaseState):
+        counter: int = 0
+
+    # AND an API node that uses the StateType parameter
+    class TestAPINode(APINode[TestState]):
+        method = APIRequestMethod.GET
+        authorization_type = AuthorizationType.BEARER_TOKEN
+        url = "https://example.com"
+        bearer_token_value = VellumSecret(name="secret")
+
+    # AND a mock response from the API
+    vellum_client.execute_api.return_value = ExecuteApiResponse(
+        status_code=200,
+        text='{"result": "success"}',
+        json_={"result": "success"},
+        headers={"content-type": "application/json"},
+    )
+
+    # WHEN we create and run the node with a state instance
+    state = TestState()
+    node = TestAPINode(state=state)
+    outputs = node.run()
+
+    # THEN the node should run successfully
+    assert vellum_client.execute_api.call_count == 1
+    assert outputs.status_code == 200
+    assert outputs.json == {"result": "success"}
+
+    # AND the state should be accessible (though not modified by the API node)
+    assert state.counter == 0

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,79 +22,130 @@ 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]
+        json: Json
         headers: Dict[str, str]
         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