"""Type definitions for MLflow webhook payloads.
This module contains class definitions for all webhook event payloads
that are sent when various model registry events occur.
"""
from typing import Literal, TypeAlias, TypedDict
from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent
[docs]class RegisteredModelCreatedPayload(TypedDict):
"""Payload sent when a new registered model is created.
Example payload:
.. code-block:: python
{
"name": "example_model",
"tags": {"example_key": "example_value"},
"description": "An example registered model",
}
"""
name: str
"""The name of the registered model."""
tags: dict[str, str]
"""Tags associated with the registered model."""
description: str | None
"""Description of the registered model."""
[docs] @classmethod
def example(cls) -> "RegisteredModelCreatedPayload":
return cls(
name="example_model",
tags={"example_key": "example_value"},
description="An example registered model",
)
[docs]class ModelVersionCreatedPayload(TypedDict):
"""Payload sent when a new model version is created.
Example payload:
.. code-block:: python
{
"name": "example_model",
"version": "1",
"source": "models:/123",
"run_id": "abcd1234abcd5678",
"tags": {"example_key": "example_value"},
"description": "An example model version",
}
"""
name: str
"""The name of the registered model."""
version: str
"""The version of the model."""
source: str
"""The source URI of the model version."""
run_id: str | None
"""The run ID associated with the model version, if applicable."""
tags: dict[str, str]
"""Tags associated with the model version."""
description: str | None
"""Description of the model version."""
[docs] @classmethod
def example(cls) -> "ModelVersionCreatedPayload":
return cls(
name="example_model",
version="1",
source="models:/123",
run_id="abcd1234abcd5678",
tags={"example_key": "example_value"},
description="An example model version",
)
[docs]class ModelVersionTagDeletedPayload(TypedDict):
"""Payload sent when a tag is deleted from a model version.
Example payload:
.. code-block:: python
{
"name": "example_model",
"version": "1",
"key": "example_key",
}
"""
name: str
"""The name of the registered model."""
version: str
"""The version of the model."""
key: str
"""The tag key being deleted."""
[docs] @classmethod
def example(cls) -> "ModelVersionTagDeletedPayload":
return cls(
name="example_model",
version="1",
key="example_key",
)
[docs]class ModelVersionAliasCreatedPayload(TypedDict):
"""
Payload sent when an alias is created for a model version.
Example payload:
.. code-block:: python
{
"name": "example_model",
"alias": "example_alias",
"version": "1",
}
"""
name: str
"""The name of the registered model."""
alias: str
"""The alias being created."""
version: str
"""The version of the model the alias is being assigned to."""
[docs] @classmethod
def example(cls) -> "ModelVersionAliasCreatedPayload":
return cls(
name="example_model",
alias="example_alias",
version="1",
)
[docs]class ModelVersionAliasDeletedPayload(TypedDict):
"""Payload sent when an alias is deleted from a model version.
Example payload:
.. code-block:: python
{
"name": "example_model",
"alias": "example_alias",
}
"""
name: str
"""The name of the registered model."""
alias: str
"""The alias being deleted."""
[docs] @classmethod
def example(cls) -> "ModelVersionAliasDeletedPayload":
return cls(
name="example_model",
alias="example_alias",
)
[docs]class PromptCreatedPayload(TypedDict):
"""Payload sent when a new prompt is created.
Example payload:
.. code-block:: python
{
"name": "example_prompt",
"tags": {"example_key": "example_value"},
"description": "An example prompt",
}
"""
name: str
"""The name of the prompt."""
tags: dict[str, str]
"""Tags associated with the prompt."""
description: str | None
"""Description of the prompt."""
[docs] @classmethod
def example(cls) -> "PromptCreatedPayload":
return cls(
name="example_prompt",
tags={"example_key": "example_value"},
description="An example prompt",
)
[docs]class PromptVersionCreatedPayload(TypedDict):
"""Payload sent when a new prompt version is created.
Example payload:
.. code-block:: python
{
"name": "example_prompt",
"version": "1",
"template": "Hello {{name}}!",
"tags": {"example_key": "example_value"},
"description": "An example prompt version",
}
"""
name: str
"""The name of the prompt."""
version: str
"""The version of the prompt."""
template: str
"""The template content of the prompt version."""
tags: dict[str, str]
"""Tags associated with the prompt version."""
description: str | None
"""Description of the prompt version."""
[docs] @classmethod
def example(cls) -> "PromptVersionCreatedPayload":
return cls(
name="example_prompt",
version="1",
template="Hello {{name}}!",
tags={"example_key": "example_value"},
description="An example prompt version",
)
[docs]class PromptTagDeletedPayload(TypedDict):
"""Payload sent when a tag is deleted from a prompt.
Example payload:
.. code-block:: python
{
"name": "example_prompt",
"key": "example_key",
}
"""
name: str
"""The name of the prompt."""
key: str
"""The tag key being deleted."""
[docs] @classmethod
def example(cls) -> "PromptTagDeletedPayload":
return cls(
name="example_prompt",
key="example_key",
)
[docs]class PromptVersionTagDeletedPayload(TypedDict):
"""Payload sent when a tag is deleted from a prompt version.
Example payload:
.. code-block:: python
{
"name": "example_prompt",
"version": "1",
"key": "example_key",
}
"""
name: str
"""The name of the prompt."""
version: str
"""The version of the prompt."""
key: str
"""The tag key being deleted."""
[docs] @classmethod
def example(cls) -> "PromptVersionTagDeletedPayload":
return cls(
name="example_prompt",
version="1",
key="example_key",
)
[docs]class PromptAliasCreatedPayload(TypedDict):
"""Payload sent when an alias is created for a prompt version.
Example payload:
.. code-block:: python
{
"name": "example_prompt",
"alias": "example_alias",
"version": "1",
}
"""
name: str
"""The name of the prompt."""
alias: str
"""The alias being created."""
version: str
"""The version of the prompt the alias is being assigned to."""
[docs] @classmethod
def example(cls) -> "PromptAliasCreatedPayload":
return cls(
name="example_prompt",
alias="example_alias",
version="1",
)
[docs]class PromptAliasDeletedPayload(TypedDict):
"""Payload sent when an alias is deleted from a prompt.
Example payload:
.. code-block:: python
{
"name": "example_prompt",
"alias": "example_alias",
}
"""
name: str
"""The name of the prompt."""
alias: str
"""The alias being deleted."""
[docs] @classmethod
def example(cls) -> "PromptAliasDeletedPayload":
return cls(
name="example_prompt",
alias="example_alias",
)
class BudgetPolicyExceededPayload(TypedDict):
"""Payload sent when a budget policy limit is exceeded.
Example payload:
.. code-block:: python
{
"budget_policy_id": "bp-abc123",
"budget_unit": "USD",
"budget_amount": 100.0,
"current_spend": 105.50,
"duration_unit": "MONTHS",
"duration_value": 1,
"target_scope": "WORKSPACE",
"workspace": "default",
"window_start": 1704067200000,
}
"""
budget_policy_id: str
"""The unique identifier of the budget policy."""
budget_unit: Literal["USD"]
"""The budget measurement unit (e.g. USD)."""
budget_amount: float
"""The budget limit amount."""
current_spend: float
"""The current cumulative spend when the limit was exceeded."""
duration_unit: Literal["MINUTES", "HOURS", "DAYS", "MONTHS"]
"""The duration unit (MINUTES, HOURS, DAYS, MONTHS)."""
duration_value: int
"""The duration value."""
target_scope: Literal["GLOBAL", "WORKSPACE"]
"""The target scope (GLOBAL or WORKSPACE)."""
workspace: str
"""The workspace this budget applies to."""
window_start: int
"""The start timestamp (milliseconds) of the current budget window."""
@classmethod
def example(cls) -> "BudgetPolicyExceededPayload":
return cls(
budget_policy_id="bp-abc123",
budget_unit="USD",
budget_amount=100.0,
current_spend=105.50,
duration_unit="MONTHS",
duration_value=1,
target_scope="WORKSPACE",
workspace="default",
window_start=1704067200000,
)
WebhookPayload: TypeAlias = (
RegisteredModelCreatedPayload
| ModelVersionCreatedPayload
| ModelVersionTagSetPayload
| ModelVersionTagDeletedPayload
| ModelVersionAliasCreatedPayload
| ModelVersionAliasDeletedPayload
| PromptCreatedPayload
| PromptVersionCreatedPayload
| PromptTagSetPayload
| PromptTagDeletedPayload
| PromptVersionTagSetPayload
| PromptVersionTagDeletedPayload
| PromptAliasCreatedPayload
| PromptAliasDeletedPayload
| BudgetPolicyExceededPayload
)
# Mapping of (entity, action) tuples to their corresponding payload classes
EVENT_TO_PAYLOAD_CLASS: dict[tuple[WebhookEntity, WebhookAction], type[WebhookPayload]] = {
(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED): RegisteredModelCreatedPayload,
(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED): ModelVersionCreatedPayload,
(WebhookEntity.MODEL_VERSION_TAG, WebhookAction.SET): ModelVersionTagSetPayload,
(WebhookEntity.MODEL_VERSION_TAG, WebhookAction.DELETED): ModelVersionTagDeletedPayload,
(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.CREATED): ModelVersionAliasCreatedPayload,
(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.DELETED): ModelVersionAliasDeletedPayload,
(WebhookEntity.PROMPT, WebhookAction.CREATED): PromptCreatedPayload,
(WebhookEntity.PROMPT_VERSION, WebhookAction.CREATED): PromptVersionCreatedPayload,
(WebhookEntity.PROMPT_TAG, WebhookAction.SET): PromptTagSetPayload,
(WebhookEntity.PROMPT_TAG, WebhookAction.DELETED): PromptTagDeletedPayload,
(WebhookEntity.PROMPT_VERSION_TAG, WebhookAction.SET): PromptVersionTagSetPayload,
(WebhookEntity.PROMPT_VERSION_TAG, WebhookAction.DELETED): PromptVersionTagDeletedPayload,
(WebhookEntity.PROMPT_ALIAS, WebhookAction.CREATED): PromptAliasCreatedPayload,
(WebhookEntity.PROMPT_ALIAS, WebhookAction.DELETED): PromptAliasDeletedPayload,
(WebhookEntity.BUDGET_POLICY, WebhookAction.EXCEEDED): BudgetPolicyExceededPayload,
}
def get_example_payload_for_event(event: WebhookEvent) -> WebhookPayload:
"""Get an example payload for the given webhook event type.
Args:
event: The webhook event instance
Returns:
Example payload for the event type
Raises:
ValueError: If the event type is unknown
"""
event_key = (event.entity, event.action)
if payload_class := EVENT_TO_PAYLOAD_CLASS.get(event_key):
return payload_class.example()
raise ValueError(f"Unknown event type: {event.entity}.{event.action}")
def get_payload_class_for_event(event: WebhookEvent) -> type[WebhookPayload] | None:
"""Get the payload class for the given webhook event type.
Args:
event: The webhook event instance
Returns:
Payload class for the event type, or None if unknown
"""
return EVENT_TO_PAYLOAD_CLASS.get((event.entity, event.action))