1  """Type definitions for MLflow webhook payloads.
  2  
  3  This module contains class definitions for all webhook event payloads
  4  that are sent when various model registry events occur.
  5  """
  6  
  7  from typing import Literal, TypeAlias, TypedDict
  8  
  9  from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent
 10  
 11  
 12  class RegisteredModelCreatedPayload(TypedDict):
 13      """Payload sent when a new registered model is created.
 14  
 15      Example payload:
 16  
 17      .. code-block:: python
 18  
 19          {
 20              "name": "example_model",
 21              "tags": {"example_key": "example_value"},
 22              "description": "An example registered model",
 23          }
 24  
 25      """
 26  
 27      name: str
 28      """The name of the registered model."""
 29      tags: dict[str, str]
 30      """Tags associated with the registered model."""
 31      description: str | None
 32      """Description of the registered model."""
 33  
 34      @classmethod
 35      def example(cls) -> "RegisteredModelCreatedPayload":
 36          return cls(
 37              name="example_model",
 38              tags={"example_key": "example_value"},
 39              description="An example registered model",
 40          )
 41  
 42  
 43  class ModelVersionCreatedPayload(TypedDict):
 44      """Payload sent when a new model version is created.
 45  
 46      Example payload:
 47  
 48      .. code-block:: python
 49  
 50          {
 51              "name": "example_model",
 52              "version": "1",
 53              "source": "models:/123",
 54              "run_id": "abcd1234abcd5678",
 55              "tags": {"example_key": "example_value"},
 56              "description": "An example model version",
 57          }
 58  
 59      """
 60  
 61      name: str
 62      """The name of the registered model."""
 63      version: str
 64      """The version of the model."""
 65      source: str
 66      """The source URI of the model version."""
 67      run_id: str | None
 68      """The run ID associated with the model version, if applicable."""
 69      tags: dict[str, str]
 70      """Tags associated with the model version."""
 71      description: str | None
 72      """Description of the model version."""
 73  
 74      @classmethod
 75      def example(cls) -> "ModelVersionCreatedPayload":
 76          return cls(
 77              name="example_model",
 78              version="1",
 79              source="models:/123",
 80              run_id="abcd1234abcd5678",
 81              tags={"example_key": "example_value"},
 82              description="An example model version",
 83          )
 84  
 85  
 86  class ModelVersionTagSetPayload(TypedDict):
 87      """Payload sent when a tag is set on a model version.
 88  
 89      Example payload:
 90  
 91      .. code-block:: python
 92  
 93          {
 94              "name": "example_model",
 95              "version": "1",
 96              "key": "example_key",
 97              "value": "example_value",
 98          }
 99  
100      """
101  
102      name: str
103      """The name of the registered model."""
104      version: str
105      """The version of the model."""
106      key: str
107      """The tag key being set."""
108      value: str
109      """The tag value being set."""
110  
111      @classmethod
112      def example(cls) -> "ModelVersionTagSetPayload":
113          return cls(
114              name="example_model",
115              version="1",
116              key="example_key",
117              value="example_value",
118          )
119  
120  
121  class ModelVersionTagDeletedPayload(TypedDict):
122      """Payload sent when a tag is deleted from a model version.
123  
124      Example payload:
125  
126      .. code-block:: python
127  
128          {
129              "name": "example_model",
130              "version": "1",
131              "key": "example_key",
132          }
133  
134      """
135  
136      name: str
137      """The name of the registered model."""
138      version: str
139      """The version of the model."""
140      key: str
141      """The tag key being deleted."""
142  
143      @classmethod
144      def example(cls) -> "ModelVersionTagDeletedPayload":
145          return cls(
146              name="example_model",
147              version="1",
148              key="example_key",
149          )
150  
151  
152  class ModelVersionAliasCreatedPayload(TypedDict):
153      """
154      Payload sent when an alias is created for a model version.
155  
156      Example payload:
157  
158      .. code-block:: python
159  
160          {
161              "name": "example_model",
162              "alias": "example_alias",
163              "version": "1",
164          }
165  
166      """
167  
168      name: str
169      """The name of the registered model."""
170      alias: str
171      """The alias being created."""
172      version: str
173      """The version of the model the alias is being assigned to."""
174  
175      @classmethod
176      def example(cls) -> "ModelVersionAliasCreatedPayload":
177          return cls(
178              name="example_model",
179              alias="example_alias",
180              version="1",
181          )
182  
183  
184  class ModelVersionAliasDeletedPayload(TypedDict):
185      """Payload sent when an alias is deleted from a model version.
186  
187      Example payload:
188  
189      .. code-block:: python
190  
191          {
192              "name": "example_model",
193              "alias": "example_alias",
194          }
195  
196      """
197  
198      name: str
199      """The name of the registered model."""
200      alias: str
201      """The alias being deleted."""
202  
203      @classmethod
204      def example(cls) -> "ModelVersionAliasDeletedPayload":
205          return cls(
206              name="example_model",
207              alias="example_alias",
208          )
209  
210  
211  class PromptCreatedPayload(TypedDict):
212      """Payload sent when a new prompt is created.
213  
214      Example payload:
215  
216      .. code-block:: python
217  
218          {
219              "name": "example_prompt",
220              "tags": {"example_key": "example_value"},
221              "description": "An example prompt",
222          }
223  
224      """
225  
226      name: str
227      """The name of the prompt."""
228      tags: dict[str, str]
229      """Tags associated with the prompt."""
230      description: str | None
231      """Description of the prompt."""
232  
233      @classmethod
234      def example(cls) -> "PromptCreatedPayload":
235          return cls(
236              name="example_prompt",
237              tags={"example_key": "example_value"},
238              description="An example prompt",
239          )
240  
241  
242  class PromptVersionCreatedPayload(TypedDict):
243      """Payload sent when a new prompt version is created.
244  
245      Example payload:
246  
247      .. code-block:: python
248  
249          {
250              "name": "example_prompt",
251              "version": "1",
252              "template": "Hello {{name}}!",
253              "tags": {"example_key": "example_value"},
254              "description": "An example prompt version",
255          }
256  
257      """
258  
259      name: str
260      """The name of the prompt."""
261      version: str
262      """The version of the prompt."""
263      template: str
264      """The template content of the prompt version."""
265      tags: dict[str, str]
266      """Tags associated with the prompt version."""
267      description: str | None
268      """Description of the prompt version."""
269  
270      @classmethod
271      def example(cls) -> "PromptVersionCreatedPayload":
272          return cls(
273              name="example_prompt",
274              version="1",
275              template="Hello {{name}}!",
276              tags={"example_key": "example_value"},
277              description="An example prompt version",
278          )
279  
280  
281  class PromptTagSetPayload(TypedDict):
282      """Payload sent when a tag is set on a prompt.
283  
284      Example payload:
285  
286      .. code-block:: python
287  
288          {
289              "name": "example_prompt",
290              "key": "example_key",
291              "value": "example_value",
292          }
293  
294      """
295  
296      name: str
297      """The name of the prompt."""
298      key: str
299      """The tag key being set."""
300      value: str
301      """The tag value being set."""
302  
303      @classmethod
304      def example(cls) -> "PromptTagSetPayload":
305          return cls(
306              name="example_prompt",
307              key="example_key",
308              value="example_value",
309          )
310  
311  
312  class PromptTagDeletedPayload(TypedDict):
313      """Payload sent when a tag is deleted from a prompt.
314  
315      Example payload:
316  
317      .. code-block:: python
318  
319          {
320              "name": "example_prompt",
321              "key": "example_key",
322          }
323  
324      """
325  
326      name: str
327      """The name of the prompt."""
328      key: str
329      """The tag key being deleted."""
330  
331      @classmethod
332      def example(cls) -> "PromptTagDeletedPayload":
333          return cls(
334              name="example_prompt",
335              key="example_key",
336          )
337  
338  
339  class PromptVersionTagSetPayload(TypedDict):
340      """Payload sent when a tag is set on a prompt version.
341  
342      Example payload:
343  
344      .. code-block:: python
345  
346          {
347              "name": "example_prompt",
348              "version": "1",
349              "key": "example_key",
350              "value": "example_value",
351          }
352  
353      """
354  
355      name: str
356      """The name of the prompt."""
357      version: str
358      """The version of the prompt."""
359      key: str
360      """The tag key being set."""
361      value: str
362      """The tag value being set."""
363  
364      @classmethod
365      def example(cls) -> "PromptVersionTagSetPayload":
366          return cls(
367              name="example_prompt",
368              version="1",
369              key="example_key",
370              value="example_value",
371          )
372  
373  
374  class PromptVersionTagDeletedPayload(TypedDict):
375      """Payload sent when a tag is deleted from a prompt version.
376  
377      Example payload:
378  
379      .. code-block:: python
380  
381          {
382              "name": "example_prompt",
383              "version": "1",
384              "key": "example_key",
385          }
386  
387      """
388  
389      name: str
390      """The name of the prompt."""
391      version: str
392      """The version of the prompt."""
393      key: str
394      """The tag key being deleted."""
395  
396      @classmethod
397      def example(cls) -> "PromptVersionTagDeletedPayload":
398          return cls(
399              name="example_prompt",
400              version="1",
401              key="example_key",
402          )
403  
404  
405  class PromptAliasCreatedPayload(TypedDict):
406      """Payload sent when an alias is created for a prompt version.
407  
408      Example payload:
409  
410      .. code-block:: python
411  
412          {
413              "name": "example_prompt",
414              "alias": "example_alias",
415              "version": "1",
416          }
417  
418      """
419  
420      name: str
421      """The name of the prompt."""
422      alias: str
423      """The alias being created."""
424      version: str
425      """The version of the prompt the alias is being assigned to."""
426  
427      @classmethod
428      def example(cls) -> "PromptAliasCreatedPayload":
429          return cls(
430              name="example_prompt",
431              alias="example_alias",
432              version="1",
433          )
434  
435  
436  class PromptAliasDeletedPayload(TypedDict):
437      """Payload sent when an alias is deleted from a prompt.
438  
439      Example payload:
440  
441      .. code-block:: python
442  
443          {
444              "name": "example_prompt",
445              "alias": "example_alias",
446          }
447  
448      """
449  
450      name: str
451      """The name of the prompt."""
452      alias: str
453      """The alias being deleted."""
454  
455      @classmethod
456      def example(cls) -> "PromptAliasDeletedPayload":
457          return cls(
458              name="example_prompt",
459              alias="example_alias",
460          )
461  
462  
463  class BudgetPolicyExceededPayload(TypedDict):
464      """Payload sent when a budget policy limit is exceeded.
465  
466      Example payload:
467  
468      .. code-block:: python
469  
470          {
471              "budget_policy_id": "bp-abc123",
472              "budget_unit": "USD",
473              "budget_amount": 100.0,
474              "current_spend": 105.50,
475              "duration_unit": "MONTHS",
476              "duration_value": 1,
477              "target_scope": "WORKSPACE",
478              "workspace": "default",
479              "window_start": 1704067200000,
480          }
481  
482      """
483  
484      budget_policy_id: str
485      """The unique identifier of the budget policy."""
486      budget_unit: Literal["USD"]
487      """The budget measurement unit (e.g. USD)."""
488      budget_amount: float
489      """The budget limit amount."""
490      current_spend: float
491      """The current cumulative spend when the limit was exceeded."""
492      duration_unit: Literal["MINUTES", "HOURS", "DAYS", "MONTHS"]
493      """The duration unit (MINUTES, HOURS, DAYS, MONTHS)."""
494      duration_value: int
495      """The duration value."""
496      target_scope: Literal["GLOBAL", "WORKSPACE"]
497      """The target scope (GLOBAL or WORKSPACE)."""
498      workspace: str
499      """The workspace this budget applies to."""
500      window_start: int
501      """The start timestamp (milliseconds) of the current budget window."""
502  
503      @classmethod
504      def example(cls) -> "BudgetPolicyExceededPayload":
505          return cls(
506              budget_policy_id="bp-abc123",
507              budget_unit="USD",
508              budget_amount=100.0,
509              current_spend=105.50,
510              duration_unit="MONTHS",
511              duration_value=1,
512              target_scope="WORKSPACE",
513              workspace="default",
514              window_start=1704067200000,
515          )
516  
517  
518  WebhookPayload: TypeAlias = (
519      RegisteredModelCreatedPayload
520      | ModelVersionCreatedPayload
521      | ModelVersionTagSetPayload
522      | ModelVersionTagDeletedPayload
523      | ModelVersionAliasCreatedPayload
524      | ModelVersionAliasDeletedPayload
525      | PromptCreatedPayload
526      | PromptVersionCreatedPayload
527      | PromptTagSetPayload
528      | PromptTagDeletedPayload
529      | PromptVersionTagSetPayload
530      | PromptVersionTagDeletedPayload
531      | PromptAliasCreatedPayload
532      | PromptAliasDeletedPayload
533      | BudgetPolicyExceededPayload
534  )
535  
536  # Mapping of (entity, action) tuples to their corresponding payload classes
537  EVENT_TO_PAYLOAD_CLASS: dict[tuple[WebhookEntity, WebhookAction], type[WebhookPayload]] = {
538      (WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED): RegisteredModelCreatedPayload,
539      (WebhookEntity.MODEL_VERSION, WebhookAction.CREATED): ModelVersionCreatedPayload,
540      (WebhookEntity.MODEL_VERSION_TAG, WebhookAction.SET): ModelVersionTagSetPayload,
541      (WebhookEntity.MODEL_VERSION_TAG, WebhookAction.DELETED): ModelVersionTagDeletedPayload,
542      (WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.CREATED): ModelVersionAliasCreatedPayload,
543      (WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.DELETED): ModelVersionAliasDeletedPayload,
544      (WebhookEntity.PROMPT, WebhookAction.CREATED): PromptCreatedPayload,
545      (WebhookEntity.PROMPT_VERSION, WebhookAction.CREATED): PromptVersionCreatedPayload,
546      (WebhookEntity.PROMPT_TAG, WebhookAction.SET): PromptTagSetPayload,
547      (WebhookEntity.PROMPT_TAG, WebhookAction.DELETED): PromptTagDeletedPayload,
548      (WebhookEntity.PROMPT_VERSION_TAG, WebhookAction.SET): PromptVersionTagSetPayload,
549      (WebhookEntity.PROMPT_VERSION_TAG, WebhookAction.DELETED): PromptVersionTagDeletedPayload,
550      (WebhookEntity.PROMPT_ALIAS, WebhookAction.CREATED): PromptAliasCreatedPayload,
551      (WebhookEntity.PROMPT_ALIAS, WebhookAction.DELETED): PromptAliasDeletedPayload,
552      (WebhookEntity.BUDGET_POLICY, WebhookAction.EXCEEDED): BudgetPolicyExceededPayload,
553  }
554  
555  
556  def get_example_payload_for_event(event: WebhookEvent) -> WebhookPayload:
557      """Get an example payload for the given webhook event type.
558  
559      Args:
560          event: The webhook event instance
561  
562      Returns:
563          Example payload for the event type
564  
565      Raises:
566          ValueError: If the event type is unknown
567      """
568      event_key = (event.entity, event.action)
569      if payload_class := EVENT_TO_PAYLOAD_CLASS.get(event_key):
570          return payload_class.example()
571  
572      raise ValueError(f"Unknown event type: {event.entity}.{event.action}")
573  
574  
575  def get_payload_class_for_event(event: WebhookEvent) -> type[WebhookPayload] | None:
576      """Get the payload class for the given webhook event type.
577  
578      Args:
579          event: The webhook event instance
580  
581      Returns:
582          Payload class for the event type, or None if unknown
583      """
584      return EVENT_TO_PAYLOAD_CLASS.get((event.entity, event.action))