vellum-ai 0.14.37__py3-none-any.whl → 0.14.39__py3-none-any.whl

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

@@ -0,0 +1,50 @@
+from vellum import TestSuiteRunMetricNumberOutput
+from vellum.client.types.chat_history_input import ChatHistoryInput
+from vellum.client.types.chat_message import ChatMessage
+from vellum.client.types.json_input import JsonInput
+from vellum.client.types.metric_definition_execution import MetricDefinitionExecution
+from vellum.client.types.number_input import NumberInput
+from vellum.client.types.string_input import StringInput
+from vellum.workflows.nodes.displayable.guardrail_node.node import GuardrailNode
+
+
+def test_guardrail_node__inputs(vellum_client):
+    """Test that GuardrailNode correctly handles inputs."""
+
+    # GIVEN a Guardrail Node with inputs
+    class MyGuard(GuardrailNode):
+        metric_definition = "example_metric_definition"
+        metric_inputs = {
+            "a_string": "hello",
+            "a_chat_history": [ChatMessage(role="USER", text="Hello, how are you?")],
+            "a_dict": {"foo": "bar"},
+            "a_int": 42,
+            "a_float": 3.14,
+        }
+
+    vellum_client.metric_definitions.execute_metric_definition.return_value = MetricDefinitionExecution(
+        outputs=[
+            TestSuiteRunMetricNumberOutput(
+                name="score",
+                value=1.0,
+            ),
+        ],
+    )
+
+    # WHEN the node is run
+    MyGuard().run()
+
+    # THEN the metric_definitions.execute_metric_definition method should be called with the correct inputs
+    mock_api = vellum_client.metric_definitions.execute_metric_definition
+    assert mock_api.call_count == 1
+
+    assert mock_api.call_args.kwargs["inputs"] == [
+        StringInput(name="a_string", type="STRING", value="hello"),
+        ChatHistoryInput(
+            name="a_chat_history", type="CHAT_HISTORY", value=[ChatMessage(role="USER", text="Hello, how are you?")]
+        ),
+        JsonInput(name="a_dict", type="JSON", value={"foo": "bar"}),
+        NumberInput(name="a_int", type="NUMBER", value=42.0),
+        NumberInput(name="a_float", type="NUMBER", value=3.14),
+    ]
+    assert len(mock_api.call_args.kwargs["inputs"]) == 5

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

@@ -170,8 +170,13 @@ def test_inline_prompt_node__function_definitions(vellum_adhoc_prompt_client):
             WorkflowErrorCode.INTERNAL_ERROR,
             "Failed to execute Prompt",
         ),
+        (
+            ApiError(status_code=403, body={"detail": "Provider credentials is missing or unavailable"}),
+            WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+            "Provider credentials is missing or unavailable",
+        ),
     ],
-    ids=["404", "invalid_dict", "invalid_body", "no_status_code", "500"],
+    ids=["404", "invalid_dict", "invalid_body", "no_status_code", "500", "403"],
 )
 def test_inline_prompt_node__api_error__invalid_inputs_node_exception(
     vellum_adhoc_prompt_client, exception, expected_code, expected_message

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

@@ -5,6 +5,8 @@ from typing import Any, Iterator, List
 
 from httpx import Response
 
+from vellum import RejectedExecutePromptEvent
+from vellum.client import ApiError
 from vellum.client.types.chat_history_input_request import ChatHistoryInputRequest
 from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.chat_message_request import ChatMessageRequest
@@ -15,6 +17,8 @@ from vellum.client.types.json_input_request import JsonInputRequest
 from vellum.client.types.prompt_output import PromptOutput
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows.context import execution_context
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.prompt_deployment_node.node import PromptDeploymentNode
 
 
@@ -194,3 +198,322 @@ def test_prompt_deployment_node__json_output(vellum_client):
     json_output = outputs[2]
     assert json_output.name == "json"
     assert json_output.value == expected_json
+
+
+def test_prompt_deployment_node__all_fallbacks_fail(vellum_client):
+    # GIVEN a Prompt Deployment Node with fallback models
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {"query": "test query"}
+        ml_model_fallbacks = ["fallback_model_1", "fallback_model_2"]
+
+    # AND all models fail with 404 errors
+    primary_error = ApiError(
+        body={"detail": "Failed to find model 'primary_model'"},
+        status_code=404,
+    )
+    fallback1_error = ApiError(
+        body={"detail": "Failed to find model 'fallback_model_1'"},
+        status_code=404,
+    )
+    fallback2_error = ApiError(
+        body={"detail": "Failed to find model 'fallback_model_2'"},
+        status_code=404,
+    )
+
+    vellum_client.execute_prompt_stream.side_effect = [primary_error, fallback1_error, fallback2_error]
+
+    # WHEN we run the node
+    node = TestPromptDeploymentNode()
+
+    # THEN an exception should be raised
+    with pytest.raises(NodeException) as exc_info:
+        list(node.run())
+
+    # AND the client should have been called three times
+    assert vellum_client.execute_prompt_stream.call_count == 3
+
+    # AND we get the expected error message
+    assert (
+        exc_info.value.message
+        == "Failed to execute prompts with these fallbacks: ['fallback_model_1', 'fallback_model_2']"
+    )
+
+
+def test_prompt_deployment_node__fallback_success(vellum_client):
+    # GIVEN a Prompt Deployment Node with fallback models
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {"query": "test query"}
+        ml_model_fallbacks = ["fallback_model_1", "fallback_model_2"]
+
+    # AND the primary model fails with a 404 error
+    primary_error = ApiError(
+        body={"detail": "Failed to find model 'primary_model'"},
+        status_code=404,
+    )
+
+    # AND the first fallback model succeeds
+    def generate_successful_stream():
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id, outputs=[StringVellumValue(value="Fallback response")]
+            ),
+        ]
+        return iter(events)
+
+    # Set up the mock to fail on primary but succeed on first fallback
+    vellum_client.execute_prompt_stream.side_effect = [primary_error, generate_successful_stream()]
+
+    # WHEN we run the node
+    node = TestPromptDeploymentNode()
+    outputs = list(node.run())
+
+    # THEN the node should complete successfully using the fallback model
+    assert len(outputs) > 0
+    assert outputs[-1].value == "Fallback response"
+
+    # AND the client should have been called twice (once for primary, once for fallback)
+    assert vellum_client.execute_prompt_stream.call_count == 2
+
+    # AND the second call should include the fallback model override
+    second_call_kwargs = vellum_client.execute_prompt_stream.call_args_list[1][1]
+    body_params = second_call_kwargs["request_options"]["additional_body_parameters"]
+    assert body_params["overrides"]["ml_model_fallback"] == "fallback_model_1"
+
+
+def test_prompt_deployment_node__provider_error_with_fallbacks(vellum_client):
+    # GIVEN a Prompt Deployment Node with fallback models
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {}
+        ml_model_fallbacks = ["gpt-4o", "gemini-1.5-flash-latest"]
+
+    # AND the primary model starts but then fails with a provider error
+    def generate_primary_events():
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            RejectedExecutePromptEvent(
+                execution_id=execution_id,
+                error={
+                    "code": "PROVIDER_ERROR",
+                    "message": "The model provider encountered an error",
+                },
+            ),
+        ]
+        return iter(events)
+
+    # AND the fallback model succeeds
+    def generate_fallback_events():
+        execution_id = str(uuid4())
+        expected_outputs: List[PromptOutput] = [StringVellumValue(value="Fallback response")]
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(execution_id=execution_id, outputs=expected_outputs),
+        ]
+        return iter(events)
+
+    vellum_client.execute_prompt_stream.side_effect = [generate_primary_events(), generate_fallback_events()]
+
+    # WHEN we run the node
+    node = TestPromptDeploymentNode()
+    outputs = list(node.run())
+
+    # THEN the node should complete successfully using the fallback model
+    assert len(outputs) > 0
+    assert outputs[-1].value == "Fallback response"
+
+    # AND the client should have been called twice
+    assert vellum_client.execute_prompt_stream.call_count == 2
+
+    # AND the second call should include the fallback model override
+    second_call_kwargs = vellum_client.execute_prompt_stream.call_args_list[1][1]
+    body_params = second_call_kwargs["request_options"]["additional_body_parameters"]
+    assert body_params["overrides"]["ml_model_fallback"] == "gpt-4o"
+
+
+def test_prompt_deployment_node__multiple_fallbacks_mixed_errors(vellum_client):
+    """
+    This test case is when the primary model fails with an api error and
+    the first fallback fails with a provider error
+    """
+
+    # GIVEN a Prompt Deployment Node with multiple fallback models
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {}
+        ml_model_fallbacks = ["gpt-4o", "gemini-1.5-flash-latest"]
+
+    # AND the primary model fails with an API error
+    primary_error = ApiError(
+        body={"detail": "Failed to find model 'primary_model'"},
+        status_code=404,
+    )
+
+    # AND the first fallback model fails with a provider error
+    def generate_fallback1_events():
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            RejectedExecutePromptEvent(
+                execution_id=execution_id,
+                error={
+                    "code": "PROVIDER_ERROR",
+                    "message": "The first fallback provider encountered an error",
+                },
+            ),
+        ]
+        return iter(events)
+
+    # AND the second fallback model succeeds
+    def generate_fallback2_events():
+        execution_id = str(uuid4())
+        expected_outputs: List[PromptOutput] = [StringVellumValue(value="Second fallback response")]
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(execution_id=execution_id, outputs=expected_outputs),
+        ]
+        return iter(events)
+
+    vellum_client.execute_prompt_stream.side_effect = [
+        primary_error,
+        generate_fallback1_events(),
+        generate_fallback2_events(),
+    ]
+
+    # WHEN we run the node
+    node = TestPromptDeploymentNode()
+    outputs = list(node.run())
+
+    # THEN the node should complete successfully using the second fallback model
+    assert len(outputs) > 0
+    assert outputs[-1].value == "Second fallback response"
+
+    # AND the client should have been called three times
+    assert vellum_client.execute_prompt_stream.call_count == 3
+
+    # AND the calls should include the correct model overrides
+    first_fallback_call = vellum_client.execute_prompt_stream.call_args_list[1][1]
+    first_fallback_params = first_fallback_call["request_options"]["additional_body_parameters"]
+    assert first_fallback_params["overrides"]["ml_model_fallback"] == "gpt-4o"
+
+    second_fallback_call = vellum_client.execute_prompt_stream.call_args_list[2][1]
+    second_fallback_params = second_fallback_call["request_options"]["additional_body_parameters"]
+    assert second_fallback_params["overrides"]["ml_model_fallback"] == "gemini-1.5-flash-latest"
+
+
+def test_prompt_deployment_node_multiple_provider_errors(vellum_client):
+    # GIVEN a Prompt Deployment Node with a single fallback model
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {}
+        ml_model_fallbacks = ["gpt-4o"]
+
+    # AND the primary model fails with a provider error
+    def generate_primary_events():
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            RejectedExecutePromptEvent(
+                execution_id=execution_id,
+                error={
+                    "code": "PROVIDER_ERROR",
+                    "message": "The primary provider encountered an error",
+                },
+            ),
+        ]
+        return iter(events)
+
+    # AND the fallback model also fails with a provider error
+    def generate_fallback1_events():
+        execution_id = str(uuid4())
+        events = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            RejectedExecutePromptEvent(
+                execution_id=execution_id,
+                error={
+                    "code": "PROVIDER_ERROR",
+                    "message": "The first fallback provider encountered an error",
+                },
+            ),
+        ]
+        return iter(events)
+
+    vellum_client.execute_prompt_stream.side_effect = [
+        generate_primary_events(),
+        generate_fallback1_events(),
+    ]
+
+    # WHEN we run the node
+    with pytest.raises(NodeException) as exc_info:
+        node = TestPromptDeploymentNode()
+        list(node.run())
+
+    # THEN we should get an exception
+    assert exc_info.value.message == "Failed to execute prompts with these fallbacks: ['gpt-4o']"
+
+    # AND the client should have been called two times
+    assert vellum_client.execute_prompt_stream.call_count == 2
+
+    # AND the calls should include the correct model overrides
+    first_fallback_call = vellum_client.execute_prompt_stream.call_args_list[1][1]
+    first_fallback_params = first_fallback_call["request_options"]["additional_body_parameters"]
+    assert first_fallback_params["overrides"]["ml_model_fallback"] == "gpt-4o"
+
+
+def test_prompt_deployment_node__no_fallbacks(vellum_client):
+    # GIVEN a Prompt Deployment Node with no fallback models
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {}
+
+    # AND the primary model fails with an API error
+    primary_error = ApiError(
+        body={"detail": "Failed to find model 'primary_model'"},
+        status_code=404,
+    )
+
+    vellum_client.execute_prompt_stream.side_effect = primary_error
+
+    # WHEN we run the node
+    node = TestPromptDeploymentNode()
+
+    # THEN the node should raise an exception
+    with pytest.raises(NodeException) as exc_info:
+        list(node.run())
+
+    # AND the exception should contain the original error message
+    assert exc_info.value.message == "Failed to find model 'primary_model'"
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+
+    # AND the client should have been called only once (for the primary model)
+    assert vellum_client.execute_prompt_stream.call_count == 1
+
+
+def test_prompt_deployment_node__provider_credentials_missing(vellum_client):
+    # GIVEN a Prompt Deployment Node
+    class TestPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test_deployment"
+        prompt_inputs = {}
+
+    # AND the client responds with a 403 error of provider credentials missing
+    primary_error = ApiError(
+        body={"detail": "Provider credentials is missing or unavailable"},
+        status_code=403,
+    )
+
+    vellum_client.execute_prompt_stream.side_effect = primary_error
+
+    # WHEN we run the node
+    node = TestPromptDeploymentNode()
+
+    # THEN the node should raise an exception
+    with pytest.raises(NodeException) as exc_info:
+        list(node.run())
+
+    # AND the exception should contain the original error message
+    assert exc_info.value.message == "Provider credentials is missing or unavailable"
+    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE