venv/lib/python3.11/site-packages/wandb/automations/actions.py

"""Actions that are triggered by W&B Automations."""

# ruff: noqa: UP007  # Avoid using `X | Y` for union fields, as this can cause issues with pydantic < 2.6

from __future__ import annotations

from typing import Any, Literal, Optional, Union

from pydantic import BeforeValidator, Field
from typing_extensions import Annotated, Self, get_args

from wandb._pydantic import GQLBase, GQLId, SerializedToJson, Typename

from ._generated import (
    AlertSeverity,
    GenericWebhookActionFields,
    GenericWebhookActionInput,
    NoOpActionFields,
    NoOpTriggeredActionInput,
    NotificationActionFields,
    NotificationActionInput,
    QueueJobActionFields,
)
from ._validators import (
    LenientStrEnum,
    default_if_none,
    to_input_action,
    to_saved_action,
    upper_if_str,
)
from .integrations import SlackIntegration, WebhookIntegration


# NOTE: Name shortened for readability and defined publicly for easier access
class ActionType(LenientStrEnum):
    """The type of action triggered by an automation."""

    QUEUE_JOB = "QUEUE_JOB"  # NOTE: Deprecated for creation
    NOTIFICATION = "NOTIFICATION"
    GENERIC_WEBHOOK = "GENERIC_WEBHOOK"
    NO_OP = "NO_OP"


# ------------------------------------------------------------------------------
# Saved types: for parsing response data from saved automations


# NOTE: `QueueJobActionInput` for defining a Launch job is deprecated,
# so while we allow parsing it from previously saved Automations, we deliberately
# don't currently expose it in the API for creating automations.
class SavedLaunchJobAction(QueueJobActionFields):
    action_type: Literal[ActionType.QUEUE_JOB] = ActionType.QUEUE_JOB


# FIXME: Find a better place to put these OR a better way to handle the
#   conversion from `InputAction` -> `SavedAction`.
#
# Necessary placeholder class defs for converting:
# - `SendNotification -> SavedNotificationAction`
# - `SendWebhook -> SavedWebhookAction`
#
# The "input" types (`Send{Notification,Webhook}`) will only have an `integration_id`,
# and we don't want/need to fetch the other `{Slack,Webhook}Integration` fields if
# we can avoid it.
class _SavedActionSlackIntegration(GQLBase, extra="allow"):
    typename__: Typename[Literal["SlackIntegration"]] = "SlackIntegration"
    id: GQLId


class _SavedActionWebhookIntegration(GQLBase, extra="allow"):
    typename__: Typename[Literal["GenericWebhookIntegration"]] = (
        "GenericWebhookIntegration"
    )
    id: GQLId


class SavedNotificationAction(NotificationActionFields):
    action_type: Literal[ActionType.NOTIFICATION] = ActionType.NOTIFICATION
    integration: _SavedActionSlackIntegration


class SavedWebhookAction(GenericWebhookActionFields):
    action_type: Literal[ActionType.GENERIC_WEBHOOK] = ActionType.GENERIC_WEBHOOK
    integration: _SavedActionWebhookIntegration

    # We override the type of the `requestPayload` field since the original GraphQL
    # schema (and generated class) effectively defines it as a string, when we know
    # and need to anticipate the expected structure of the JSON-serialized data.
    request_payload: Annotated[
        Optional[SerializedToJson[dict[str, Any]]],
        Field(alias="requestPayload"),
    ] = None  # type: ignore[assignment]


class SavedNoOpAction(NoOpActionFields, frozen=True):
    action_type: Literal[ActionType.NO_OP] = ActionType.NO_OP

    no_op: Annotated[bool, BeforeValidator(default_if_none)] = True
    """Placeholder field, only needed to conform to schema requirements.

    There should never be a need to set this field explicitly, as its value is ignored.
    """


# for type annotations
SavedAction = Annotated[
    Union[
        SavedLaunchJobAction,
        SavedNotificationAction,
        SavedWebhookAction,
        SavedNoOpAction,
    ],
    BeforeValidator(to_saved_action),
    Field(discriminator="typename__"),
]
# for runtime type checks
SavedActionTypes: tuple[type, ...] = get_args(SavedAction.__origin__)  # type: ignore[attr-defined]


# ------------------------------------------------------------------------------
# Input types: for creating or updating automations
class _BaseActionInput(GQLBase):
    action_type: Annotated[ActionType, Field(frozen=True)]
    """The kind of action to be triggered."""


class SendNotification(_BaseActionInput, NotificationActionInput):
    """Defines an automation action that sends a (Slack) notification."""

    action_type: Literal[ActionType.NOTIFICATION] = ActionType.NOTIFICATION

    integration_id: GQLId
    """The ID of the Slack integration that will be used to send the notification."""

    # Note: Validation aliases are meant to provide continuity with prior `wandb.alert()` API.
    title: str = ""
    """The title of the sent notification."""

    message: Annotated[str, Field(validation_alias="text")] = ""
    """The message body of the sent notification."""

    severity: Annotated[
        AlertSeverity,
        BeforeValidator(upper_if_str),  # Be helpful by ensuring uppercase strings
        Field(validation_alias="level"),
    ] = AlertSeverity.INFO
    """The severity (`INFO`, `WARN`, `ERROR`) of the sent notification."""

    @classmethod
    def from_integration(
        cls,
        integration: SlackIntegration,
        *,
        title: str = "",
        text: str = "",
        level: AlertSeverity = AlertSeverity.INFO,
    ) -> Self:
        """Define a notification action that sends to the given (Slack) integration."""
        return cls(
            integration_id=integration.id,
            title=title,
            message=text,
            severity=level,
        )


class SendWebhook(_BaseActionInput, GenericWebhookActionInput):
    """Defines an automation action that sends a webhook request."""

    action_type: Literal[ActionType.GENERIC_WEBHOOK] = ActionType.GENERIC_WEBHOOK

    integration_id: GQLId
    """The ID of the webhook integration that will be used to send the request."""

    # overrides the generated field type to parse/serialize JSON strings
    request_payload: Optional[SerializedToJson[dict[str, Any]]] = Field(  # type: ignore[assignment]
        default=None, alias="requestPayload"
    )
    """The payload, possibly with template variables, to send in the webhook request."""

    @classmethod
    def from_integration(
        cls,
        integration: WebhookIntegration,
        *,
        payload: Optional[SerializedToJson[dict[str, Any]]] = None,
    ) -> Self:
        """Define a webhook action that sends to the given (webhook) integration."""
        return cls(integration_id=integration.id, request_payload=payload)


class DoNothing(_BaseActionInput, NoOpTriggeredActionInput, frozen=True):
    """Defines an automation action that intentionally does nothing."""

    action_type: Literal[ActionType.NO_OP] = ActionType.NO_OP

    no_op: Annotated[bool, BeforeValidator(default_if_none)] = True
    """Placeholder field which exists only to satisfy backend schema requirements.

    There should never be a need to set this field explicitly, as its value is ignored.
    """


# for type annotations
InputAction = Annotated[
    Union[
        SendNotification,
        SendWebhook,
        DoNothing,
    ],
    BeforeValidator(to_input_action),
    Field(discriminator="action_type"),
]
# for runtime type checks
InputActionTypes: tuple[type, ...] = get_args(InputAction.__origin__)  # type: ignore[attr-defined]

__all__ = [
    "ActionType",
    *(cls.__name__ for cls in InputActionTypes),
]