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

vellum/client/types/__init__.py

@@ -453,6 +453,9 @@ from .test_suite_run_metric_json_output import TestSuiteRunMetricJsonOutput
 from .test_suite_run_metric_number_output import TestSuiteRunMetricNumberOutput
 from .test_suite_run_metric_output import TestSuiteRunMetricOutput
 from .test_suite_run_metric_string_output import TestSuiteRunMetricStringOutput
+from .test_suite_run_progress import TestSuiteRunProgress
+from .test_suite_run_prompt_sandbox_exec_config_data_request import TestSuiteRunPromptSandboxExecConfigDataRequest
+from .test_suite_run_prompt_sandbox_exec_config_request import TestSuiteRunPromptSandboxExecConfigRequest
 from .test_suite_run_prompt_sandbox_history_item_exec_config import TestSuiteRunPromptSandboxHistoryItemExecConfig
 from .test_suite_run_prompt_sandbox_history_item_exec_config_data import (
     TestSuiteRunPromptSandboxHistoryItemExecConfigData,
@@ -472,6 +475,8 @@ from .test_suite_run_workflow_release_tag_exec_config_data_request import (
     TestSuiteRunWorkflowReleaseTagExecConfigDataRequest,
 )
 from .test_suite_run_workflow_release_tag_exec_config_request import TestSuiteRunWorkflowReleaseTagExecConfigRequest
+from .test_suite_run_workflow_sandbox_exec_config_data_request import TestSuiteRunWorkflowSandboxExecConfigDataRequest
+from .test_suite_run_workflow_sandbox_exec_config_request import TestSuiteRunWorkflowSandboxExecConfigRequest
 from .test_suite_run_workflow_sandbox_history_item_exec_config import TestSuiteRunWorkflowSandboxHistoryItemExecConfig
 from .test_suite_run_workflow_sandbox_history_item_exec_config_data import (
     TestSuiteRunWorkflowSandboxHistoryItemExecConfigData,
@@ -1048,6 +1053,9 @@ __all__ = [
     "TestSuiteRunMetricNumberOutput",
     "TestSuiteRunMetricOutput",
     "TestSuiteRunMetricStringOutput",
+    "TestSuiteRunProgress",
+    "TestSuiteRunPromptSandboxExecConfigDataRequest",
+    "TestSuiteRunPromptSandboxExecConfigRequest",
     "TestSuiteRunPromptSandboxHistoryItemExecConfig",
     "TestSuiteRunPromptSandboxHistoryItemExecConfigData",
     "TestSuiteRunPromptSandboxHistoryItemExecConfigDataRequest",
@@ -1059,6 +1067,8 @@ __all__ = [
     "TestSuiteRunWorkflowReleaseTagExecConfigData",
     "TestSuiteRunWorkflowReleaseTagExecConfigDataRequest",
     "TestSuiteRunWorkflowReleaseTagExecConfigRequest",
+    "TestSuiteRunWorkflowSandboxExecConfigDataRequest",
+    "TestSuiteRunWorkflowSandboxExecConfigRequest",
     "TestSuiteRunWorkflowSandboxHistoryItemExecConfig",
     "TestSuiteRunWorkflowSandboxHistoryItemExecConfigData",
     "TestSuiteRunWorkflowSandboxHistoryItemExecConfigDataRequest",

vellum/client/types/ad_hoc_fulfilled_prompt_execution_meta.py

@@ -4,6 +4,7 @@ from ..core.pydantic_utilities import UniversalBaseModel
 import typing
 from .finish_reason_enum import FinishReasonEnum
 from .ml_model_usage import MlModelUsage
+from .price import Price
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
 
@@ -16,6 +17,7 @@ class AdHocFulfilledPromptExecutionMeta(UniversalBaseModel):
     latency: typing.Optional[int] = None
     finish_reason: typing.Optional[FinishReasonEnum] = None
     usage: typing.Optional[MlModelUsage] = None
+    cost: typing.Optional[Price] = None
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

vellum/client/types/fulfilled_prompt_execution_meta.py

@@ -4,6 +4,7 @@ from ..core.pydantic_utilities import UniversalBaseModel
 import typing
 from .finish_reason_enum import FinishReasonEnum
 from .ml_model_usage import MlModelUsage
+from .price import Price
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
 
@@ -16,6 +17,7 @@ class FulfilledPromptExecutionMeta(UniversalBaseModel):
     latency: typing.Optional[int] = None
     finish_reason: typing.Optional[FinishReasonEnum] = None
     usage: typing.Optional[MlModelUsage] = None
+    cost: typing.Optional[Price] = None
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

vellum/client/types/test_suite_run_exec_config_request.py

@@ -2,10 +2,12 @@
 
 import typing
 from .test_suite_run_deployment_release_tag_exec_config_request import TestSuiteRunDeploymentReleaseTagExecConfigRequest
+from .test_suite_run_prompt_sandbox_exec_config_request import TestSuiteRunPromptSandboxExecConfigRequest
 from .test_suite_run_prompt_sandbox_history_item_exec_config_request import (
     TestSuiteRunPromptSandboxHistoryItemExecConfigRequest,
 )
 from .test_suite_run_workflow_release_tag_exec_config_request import TestSuiteRunWorkflowReleaseTagExecConfigRequest
+from .test_suite_run_workflow_sandbox_exec_config_request import TestSuiteRunWorkflowSandboxExecConfigRequest
 from .test_suite_run_workflow_sandbox_history_item_exec_config_request import (
     TestSuiteRunWorkflowSandboxHistoryItemExecConfigRequest,
 )
@@ -13,8 +15,10 @@ from .test_suite_run_external_exec_config_request import TestSuiteRunExternalExe
 
 TestSuiteRunExecConfigRequest = typing.Union[
     TestSuiteRunDeploymentReleaseTagExecConfigRequest,
+    TestSuiteRunPromptSandboxExecConfigRequest,
     TestSuiteRunPromptSandboxHistoryItemExecConfigRequest,
     TestSuiteRunWorkflowReleaseTagExecConfigRequest,
+    TestSuiteRunWorkflowSandboxExecConfigRequest,
     TestSuiteRunWorkflowSandboxHistoryItemExecConfigRequest,
     TestSuiteRunExternalExecConfigRequest,
 ]

vellum/client/types/test_suite_run_progress.py

@@ -0,0 +1,20 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+import pydantic
+
+
+class TestSuiteRunProgress(UniversalBaseModel):
+    number_of_requested_test_cases: int
+    number_of_completed_test_cases: int
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/test_suite_run_prompt_sandbox_exec_config_data_request.py

@@ -0,0 +1,27 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+
+
+class TestSuiteRunPromptSandboxExecConfigDataRequest(UniversalBaseModel):
+    prompt_sandbox_id: str = pydantic.Field()
+    """
+    The ID of the Prompt Sandbox to run the Test Suite against.
+    """
+
+    prompt_variant_id: str = pydantic.Field()
+    """
+    The ID of the Prompt Variant within the Prompt Sandbox that you'd like to run the Test Suite against.
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/test_suite_run_prompt_sandbox_exec_config_request.py

@@ -0,0 +1,29 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .test_suite_run_prompt_sandbox_exec_config_data_request import TestSuiteRunPromptSandboxExecConfigDataRequest
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+
+
+class TestSuiteRunPromptSandboxExecConfigRequest(UniversalBaseModel):
+    """
+    Execution configuration for running a Test Suite against a Prompt Sandbox
+    """
+
+    type: typing.Literal["PROMPT_SANDBOX"] = "PROMPT_SANDBOX"
+    data: TestSuiteRunPromptSandboxExecConfigDataRequest
+    test_case_ids: typing.Optional[typing.List[str]] = pydantic.Field(default=None)
+    """
+    Optionally specify a subset of test case ids to run. If not provided, all test cases within the test suite will be run by default.
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/test_suite_run_read.py

@@ -9,6 +9,7 @@ from .test_suite_run_state import TestSuiteRunState
 import pydantic
 import typing
 from .test_suite_run_exec_config import TestSuiteRunExecConfig
+from .test_suite_run_progress import TestSuiteRunProgress
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 from ..core.pydantic_utilities import update_forward_refs
 
@@ -33,6 +34,8 @@ class TestSuiteRunRead(UniversalBaseModel):
     Configuration that defines how the Test Suite should be run
     """
 
+    progress: typing.Optional[TestSuiteRunProgress] = None
+
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
     else:

vellum/client/types/test_suite_run_workflow_sandbox_exec_config_data_request.py

@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+
+
+class TestSuiteRunWorkflowSandboxExecConfigDataRequest(UniversalBaseModel):
+    workflow_sandbox_id: str = pydantic.Field()
+    """
+    The ID of the Workflow Sandbox to run the Test Suite against.
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/test_suite_run_workflow_sandbox_exec_config_request.py

@@ -0,0 +1,29 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .test_suite_run_workflow_sandbox_exec_config_data_request import TestSuiteRunWorkflowSandboxExecConfigDataRequest
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+
+
+class TestSuiteRunWorkflowSandboxExecConfigRequest(UniversalBaseModel):
+    """
+    Execution configuration for running a Test Suite against a Workflow Sandbox
+    """
+
+    type: typing.Literal["WORKFLOW_SANDBOX"] = "WORKFLOW_SANDBOX"
+    data: TestSuiteRunWorkflowSandboxExecConfigDataRequest
+    test_case_ids: typing.Optional[typing.List[str]] = pydantic.Field(default=None)
+    """
+    Optionally specify a subset of test case ids to run. If not provided, all test cases within the test suite will be run by default.
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/vellum_sdk_error_code_enum.py

@@ -11,6 +11,7 @@ VellumSdkErrorCodeEnum = typing.Union[
         "INVALID_CODE",
         "INVALID_TEMPLATE",
         "INTERNAL_ERROR",
+        "PROVIDER_CREDENTIALS_UNAVAILABLE",
         "PROVIDER_ERROR",
         "USER_DEFINED_ERROR",
         "WORKFLOW_CANCELLED",

vellum/client/types/workflow_execution_event_error_code.py

@@ -6,6 +6,7 @@ WorkflowExecutionEventErrorCode = typing.Union[
     typing.Literal[
         "WORKFLOW_INITIALIZATION",
         "WORKFLOW_CANCELLED",
+        "PROVIDER_CREDENTIALS_UNAVAILABLE",
         "NODE_EXECUTION_COUNT_LIMIT_REACHED",
         "INTERNAL_SERVER_ERROR",
         "NODE_EXECUTION",

vellum/plugins/pydantic.py

@@ -53,7 +53,7 @@ class OnValidatePython(ValidatePythonHandlerProtocol):
             return
 
         if self_instance:
-            model_fields: Dict[str, FieldInfo] = self_instance.model_fields
+            model_fields: Dict[str, FieldInfo] = self_instance.__class__.model_fields
         else:
             model_fields = {}
 

vellum/types/test_suite_run_progress.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_suite_run_progress import *

vellum/types/test_suite_run_prompt_sandbox_exec_config_data_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.test_suite_run_prompt_sandbox_exec_config_data_request import *

vellum/types/test_suite_run_prompt_sandbox_exec_config_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.test_suite_run_prompt_sandbox_exec_config_request import *

vellum/types/test_suite_run_workflow_sandbox_exec_config_data_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.test_suite_run_workflow_sandbox_exec_config_data_request import *

vellum/types/test_suite_run_workflow_sandbox_exec_config_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.test_suite_run_workflow_sandbox_exec_config_request import *

vellum/workflows/errors/types.py

@@ -17,6 +17,7 @@ class WorkflowErrorCode(Enum):
     INVALID_TEMPLATE = "INVALID_TEMPLATE"
     INTERNAL_ERROR = "INTERNAL_ERROR"
     NODE_EXECUTION = "NODE_EXECUTION"
+    PROVIDER_CREDENTIALS_UNAVAILABLE = "PROVIDER_CREDENTIALS_UNAVAILABLE"
     PROVIDER_ERROR = "PROVIDER_ERROR"
     USER_DEFINED_ERROR = "USER_DEFINED_ERROR"
     WORKFLOW_CANCELLED = "WORKFLOW_CANCELLED"

vellum/workflows/events/node.py

@@ -8,9 +8,10 @@ from vellum.workflows.expressions.accessor import AccessorExpression
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references.node import NodeReference
+from vellum.workflows.types.definition import serialize_type_encoder_with_id
 from vellum.workflows.types.generics import OutputsType
 
-from .types import BaseEvent, default_serializer, serialize_type_encoder_with_id
+from .types import BaseEvent, default_serializer
 
 if TYPE_CHECKING:
     from vellum.workflows.nodes.bases import BaseNode

vellum/workflows/events/tests/test_event.py

@@ -89,6 +89,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                         "foo": "bar",
                     },
                     "display_context": None,
+                    "initial_state": None,
                 },
                 "parent": None,
             },

vellum/workflows/events/types.py

@@ -1,13 +1,14 @@
 from datetime import datetime
 import json
 from uuid import UUID, uuid4
-from typing import Annotated, Any, Dict, List, Literal, Optional, Union, get_args
+from typing import Annotated, Any, Literal, Optional, Union, get_args
 
-from pydantic import BeforeValidator, Field, GetCoreSchemaHandler, Tag, ValidationInfo
+from pydantic import Field, GetCoreSchemaHandler, Tag, ValidationInfo
 from pydantic_core import CoreSchema, core_schema
 
 from vellum.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.state.encoder import DefaultStateEncoder
+from vellum.workflows.types.definition import VellumCodeResourceDefinition
 from vellum.workflows.types.utils import datetime_now
 
 
@@ -19,28 +20,6 @@ def default_datetime_factory() -> datetime:
     return datetime_now()
 
 
-excluded_modules = {"typing", "builtins"}
-
-
-def serialize_type_encoder(obj: type) -> Dict[str, Any]:
-    return {
-        "name": obj.__name__,
-        "module": obj.__module__.split("."),
-    }
-
-
-def serialize_type_encoder_with_id(obj: Union[type, "CodeResourceDefinition"]) -> Dict[str, Any]:
-    if hasattr(obj, "__id__") and isinstance(obj, type):
-        return {
-            "id": getattr(obj, "__id__"),
-            **serialize_type_encoder(obj),
-        }
-    elif isinstance(obj, CodeResourceDefinition):
-        return obj.model_dump(mode="json")
-
-    raise AttributeError(f"The object of type '{type(obj).__name__}' must have an '__id__' attribute.")
-
-
 def default_serializer(obj: Any) -> Any:
     return json.loads(
         json.dumps(
@@ -50,22 +29,6 @@ def default_serializer(obj: Any) -> Any:
     )
 
 
-class CodeResourceDefinition(UniversalBaseModel):
-    id: UUID
-    name: str
-    module: List[str]
-
-    @staticmethod
-    def encode(obj: type) -> "CodeResourceDefinition":
-        return CodeResourceDefinition(**serialize_type_encoder_with_id(obj))
-
-
-VellumCodeResourceDefinition = Annotated[
-    CodeResourceDefinition,
-    BeforeValidator(lambda d: (d if type(d) is dict else serialize_type_encoder_with_id(d))),
-]
-
-
 class BaseParentContext(UniversalBaseModel):
     span_id: UUID
     parent: Optional["ParentContext"] = None

vellum/workflows/events/workflow.py

@@ -8,6 +8,7 @@ from vellum.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.errors import WorkflowError
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.references import ExternalInputReference
+from vellum.workflows.types.definition import serialize_type_encoder_with_id
 from vellum.workflows.types.generics import InputsType, OutputsType, StateType
 
 from .node import (
@@ -18,7 +19,7 @@ from .node import (
     NodeExecutionResumedEvent,
     NodeExecutionStreamingEvent,
 )
-from .types import BaseEvent, default_serializer, serialize_type_encoder_with_id
+from .types import BaseEvent, default_serializer
 
 if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
@@ -53,8 +54,10 @@ class WorkflowEventDisplayContext(UniversalBaseModel):
     workflow_outputs: Dict[str, UUID]
 
 
-class WorkflowExecutionInitiatedBody(_BaseWorkflowExecutionBody, Generic[InputsType]):
+class WorkflowExecutionInitiatedBody(_BaseWorkflowExecutionBody, Generic[InputsType, StateType]):
     inputs: InputsType
+    initial_state: Optional[StateType] = None
+
     # It is still the responsibility of the workflow server to populate this context. The SDK's
     # Workflow Runner will always leave this field None.
     #
@@ -66,15 +69,23 @@ class WorkflowExecutionInitiatedBody(_BaseWorkflowExecutionBody, Generic[InputsT
     def serialize_inputs(self, inputs: InputsType, _info: Any) -> Dict[str, Any]:
         return default_serializer(inputs)
 
+    @field_serializer("initial_state")
+    def serialize_initial_state(self, initial_state: Optional[StateType], _info: Any) -> Optional[Dict[str, Any]]:
+        return default_serializer(initial_state)
+
 
-class WorkflowExecutionInitiatedEvent(_BaseWorkflowEvent, Generic[InputsType]):
+class WorkflowExecutionInitiatedEvent(_BaseWorkflowEvent, Generic[InputsType, StateType]):
     name: Literal["workflow.execution.initiated"] = "workflow.execution.initiated"
-    body: WorkflowExecutionInitiatedBody[InputsType]
+    body: WorkflowExecutionInitiatedBody[InputsType, StateType]
 
     @property
     def inputs(self) -> InputsType:
         return self.body.inputs
 
+    @property
+    def initial_state(self) -> Optional[StateType]:
+        return self.body.initial_state
+
 
 class WorkflowExecutionStreamingBody(_BaseWorkflowExecutionBody):
     output: BaseOutput

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

@@ -69,7 +69,13 @@ class BasePromptNode(BaseNode, Generic[StateType]):
         return outputs
 
     def _handle_api_error(self, e: ApiError):
-        if e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
+        if e.status_code and e.status_code == 403 and isinstance(e.body, dict):
+            raise NodeException(
+                message=e.body.get("detail", "Provider credentials is missing or unavailable"),
+                code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+            )
+
+        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,

vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py

@@ -1,6 +1,6 @@
 import json
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Sequence, Union
+from typing import Any, ClassVar, Dict, Generator, Generic, Iterator, List, Optional, Sequence, Set, Union
 
 from vellum import (
     ChatHistoryInputRequest,
@@ -9,17 +9,20 @@ from vellum import (
     JsonInputRequest,
     PromptDeploymentExpandMetaRequest,
     PromptDeploymentInputRequest,
+    PromptOutput,
     RawPromptExecutionOverridesRequest,
     StringInputRequest,
 )
-from vellum.client import RequestOptions
+from vellum.client import ApiError, RequestOptions
 from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
 from vellum.workflows.context import get_execution_context
 from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.errors.types import vellum_error_to_workflow_error
 from vellum.workflows.events.types import default_serializer
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
+from vellum.workflows.outputs import BaseOutput
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 
@@ -56,13 +59,21 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
     class Trigger(BasePromptNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
-    def _get_prompt_event_stream(self) -> Iterator[ExecutePromptEvent]:
+    def _get_prompt_event_stream(self, ml_model_fallback: Optional[str] = None) -> Iterator[ExecutePromptEvent]:
         execution_context = get_execution_context()
         request_options = self.request_options or RequestOptions()
         request_options["additional_body_parameters"] = {
             "execution_context": execution_context.model_dump(mode="json"),
             **request_options.get("additional_body_parameters", {}),
         }
+        if ml_model_fallback:
+            request_options["additional_body_parameters"] = {
+                "overrides": {
+                    "ml_model_fallback": ml_model_fallback,
+                },
+                **request_options.get("additional_body_parameters", {}),
+            }
+
         return self._context.vellum_client.execute_prompt_stream(
             inputs=self._compile_prompt_inputs(),
             prompt_deployment_id=str(self.deployment) if isinstance(self.deployment, UUID) else None,
@@ -76,6 +87,86 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
             request_options=request_options,
         )
 
+    def _process_prompt_event_stream(
+        self,
+        prompt_event_stream: Optional[Iterator[ExecutePromptEvent]] = None,
+        tried_fallbacks: Optional[set[str]] = None,
+    ) -> Generator[BaseOutput, None, Optional[List[PromptOutput]]]:
+        """Override the base prompt node _process_prompt_event_stream()"""
+        self._validate()
+
+        if tried_fallbacks is None:
+            tried_fallbacks = set()
+
+        if prompt_event_stream is None:
+            try:
+                prompt_event_stream = self._get_prompt_event_stream()
+                next(prompt_event_stream)
+            except ApiError as e:
+                if (
+                    e.status_code
+                    and e.status_code < 500
+                    and self.ml_model_fallbacks is not OMIT
+                    and self.ml_model_fallbacks is not None
+                ):
+                    prompt_event_stream = self._retry_prompt_stream_with_fallbacks(tried_fallbacks)
+                else:
+                    self._handle_api_error(e)
+
+        outputs: Optional[List[PromptOutput]] = None
+        if prompt_event_stream is not None:
+            for event in prompt_event_stream:
+                if event.state == "INITIATED":
+                    continue
+                elif event.state == "STREAMING":
+                    yield BaseOutput(name="results", delta=event.output.value)
+                elif event.state == "FULFILLED":
+                    outputs = event.outputs
+                    yield BaseOutput(name="results", value=event.outputs)
+                elif event.state == "REJECTED":
+                    if (
+                        event.error
+                        and event.error.code == WorkflowErrorCode.PROVIDER_ERROR.value
+                        and self.ml_model_fallbacks is not OMIT
+                        and self.ml_model_fallbacks is not None
+                    ):
+                        try:
+                            fallback_stream = self._retry_prompt_stream_with_fallbacks(tried_fallbacks)
+                            fallback_outputs = yield from self._process_prompt_event_stream(
+                                fallback_stream, tried_fallbacks
+                            )
+                            return fallback_outputs
+                        except ApiError:
+                            pass
+
+                    workflow_error = vellum_error_to_workflow_error(event.error)
+                    raise NodeException.of(workflow_error)
+
+        return outputs
+
+    def _retry_prompt_stream_with_fallbacks(self, tried_fallbacks: Set[str]) -> Optional[Iterator[ExecutePromptEvent]]:
+        if self.ml_model_fallbacks is not None:
+            for ml_model_fallback in self.ml_model_fallbacks:
+                if ml_model_fallback in tried_fallbacks:
+                    continue
+
+                try:
+                    tried_fallbacks.add(ml_model_fallback)
+                    prompt_event_stream = self._get_prompt_event_stream(ml_model_fallback=ml_model_fallback)
+                    next(prompt_event_stream)
+                    return prompt_event_stream
+                except ApiError:
+                    continue
+            else:
+                self._handle_api_error(
+                    ApiError(
+                        body={"detail": f"Failed to execute prompts with these fallbacks: {self.ml_model_fallbacks}"},
+                        status_code=400,
+                    )
+                )
+
+        return None
+
     def _compile_prompt_inputs(self) -> List[PromptDeploymentInputRequest]:
         # TODO: We may want to consolidate with subworkflow deployment input compilation
         # https://app.shortcut.com/vellum/story/4117

vellum/workflows/nodes/displayable/conftest.py

@@ -1,12 +1,8 @@
 import pytest
 from uuid import UUID
 
-from vellum.workflows.events.types import (
-    CodeResourceDefinition,
-    NodeParentContext,
-    WorkflowDeploymentParentContext,
-    WorkflowParentContext,
-)
+from vellum.workflows.events.types import NodeParentContext, WorkflowDeploymentParentContext, WorkflowParentContext
+from vellum.workflows.types.definition import CodeResourceDefinition
 
 
 @pytest.fixture

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

@@ -117,7 +117,7 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
                         value=cast(Dict[str, Any], input_value),
                     )
                 )
-            elif isinstance(input_value, float):
+            elif isinstance(input_value, (int, float)):
                 compiled_inputs.append(
                     NumberInput(
                         name=input_name,