1  from enum import Enum
  2  from typing import Literal, TypeAlias
  3  
  4  from typing_extensions import Self
  5  
  6  from mlflow.exceptions import MlflowException
  7  from mlflow.protos.webhooks_pb2 import Webhook as ProtoWebhook
  8  from mlflow.protos.webhooks_pb2 import WebhookAction as ProtoWebhookAction
  9  from mlflow.protos.webhooks_pb2 import WebhookEntity as ProtoWebhookEntity
 10  from mlflow.protos.webhooks_pb2 import WebhookEvent as ProtoWebhookEvent
 11  from mlflow.protos.webhooks_pb2 import WebhookStatus as ProtoWebhookStatus
 12  from mlflow.protos.webhooks_pb2 import WebhookTestResult as ProtoWebhookTestResult
 13  from mlflow.utils.workspace_utils import resolve_entity_workspace_name
 14  
 15  
 16  class WebhookStatus(str, Enum):
 17      ACTIVE = "ACTIVE"
 18      DISABLED = "DISABLED"
 19  
 20      def __str__(self) -> str:
 21          return self.value
 22  
 23      @classmethod
 24      def from_proto(cls, proto: int) -> Self:
 25          proto_name = ProtoWebhookStatus.Name(proto)
 26          try:
 27              return cls(proto_name)
 28          except ValueError:
 29              raise ValueError(f"Unknown proto status: {proto_name}")
 30  
 31      def to_proto(self) -> int:
 32          return ProtoWebhookStatus.Value(self.value)
 33  
 34      def is_active(self) -> bool:
 35          return self == WebhookStatus.ACTIVE
 36  
 37  
 38  class WebhookEntity(str, Enum):
 39      REGISTERED_MODEL = "registered_model"
 40      MODEL_VERSION = "model_version"
 41      MODEL_VERSION_TAG = "model_version_tag"
 42      MODEL_VERSION_ALIAS = "model_version_alias"
 43      PROMPT = "prompt"
 44      PROMPT_VERSION = "prompt_version"
 45      PROMPT_TAG = "prompt_tag"
 46      PROMPT_VERSION_TAG = "prompt_version_tag"
 47      PROMPT_ALIAS = "prompt_alias"
 48      BUDGET_POLICY = "budget_policy"
 49  
 50      def __str__(self) -> str:
 51          return self.value
 52  
 53      @classmethod
 54      def from_proto(cls, proto: int) -> Self:
 55          proto_name = ProtoWebhookEntity.Name(proto)
 56          entity_value = proto_name.lower()
 57          return cls(entity_value)
 58  
 59      def to_proto(self) -> int:
 60          proto_name = self.value.upper()
 61          return ProtoWebhookEntity.Value(proto_name)
 62  
 63  
 64  class WebhookAction(str, Enum):
 65      CREATED = "created"
 66      UPDATED = "updated"
 67      DELETED = "deleted"
 68      SET = "set"
 69      EXCEEDED = "exceeded"
 70  
 71      def __str__(self) -> str:
 72          return self.value
 73  
 74      @classmethod
 75      def from_proto(cls, proto: int) -> Self:
 76          proto_name = ProtoWebhookAction.Name(proto)
 77          # Convert UPPER_CASE to lowercase
 78          action_value = proto_name.lower()
 79          try:
 80              return cls(action_value)
 81          except ValueError:
 82              raise ValueError(f"Unknown proto action: {proto_name}")
 83  
 84      def to_proto(self) -> int:
 85          # Convert lowercase to UPPER_CASE
 86          proto_name = self.value.upper()
 87          return ProtoWebhookAction.Value(proto_name)
 88  
 89  
 90  WebhookEventStr: TypeAlias = Literal[
 91      "registered_model.created",
 92      "model_version.created",
 93      "model_version_tag.set",
 94      "model_version_tag.deleted",
 95      "model_version_alias.created",
 96      "model_version_alias.deleted",
 97      "prompt.created",
 98      "prompt_version.created",
 99      "prompt_tag.set",
100      "prompt_tag.deleted",
101      "prompt_version_tag.set",
102      "prompt_version_tag.deleted",
103      "prompt_alias.created",
104      "prompt_alias.deleted",
105      "budget_policy.exceeded",
106  ]
107  
108  # Valid actions for each entity type
109  VALID_ENTITY_ACTIONS: dict[WebhookEntity, set[WebhookAction]] = {
110      WebhookEntity.REGISTERED_MODEL: {
111          WebhookAction.CREATED,
112      },
113      WebhookEntity.MODEL_VERSION: {
114          WebhookAction.CREATED,
115      },
116      WebhookEntity.MODEL_VERSION_TAG: {
117          WebhookAction.SET,
118          WebhookAction.DELETED,
119      },
120      WebhookEntity.MODEL_VERSION_ALIAS: {
121          WebhookAction.CREATED,
122          WebhookAction.DELETED,
123      },
124      WebhookEntity.PROMPT: {
125          WebhookAction.CREATED,
126      },
127      WebhookEntity.PROMPT_VERSION: {
128          WebhookAction.CREATED,
129      },
130      WebhookEntity.PROMPT_TAG: {
131          WebhookAction.SET,
132          WebhookAction.DELETED,
133      },
134      WebhookEntity.PROMPT_VERSION_TAG: {
135          WebhookAction.SET,
136          WebhookAction.DELETED,
137      },
138      WebhookEntity.PROMPT_ALIAS: {
139          WebhookAction.CREATED,
140          WebhookAction.DELETED,
141      },
142      WebhookEntity.BUDGET_POLICY: {
143          WebhookAction.EXCEEDED,
144      },
145  }
146  
147  
148  class WebhookEvent:
149      """
150      Represents a webhook event with a resource and action.
151      """
152  
153      def __init__(
154          self,
155          entity: str | WebhookEntity,
156          action: str | WebhookAction,
157      ):
158          """
159          Initialize a WebhookEvent.
160  
161          Args:
162              entity: The entity type (string or WebhookEntity enum)
163              action: The action type (string or WebhookAction enum)
164  
165          Raises:
166              MlflowException: If the entity/action combination is invalid
167          """
168          self._entity = WebhookEntity(entity) if isinstance(entity, str) else entity
169          self._action = WebhookAction(action) if isinstance(action, str) else action
170  
171          # Validate entity/action combination
172          if not self._is_valid_combination(self._entity, self._action):
173              valid_actions = VALID_ENTITY_ACTIONS.get(self._entity, set())
174              raise MlflowException.invalid_parameter_value(
175                  f"Invalid action '{self._action}' for entity '{self._entity}'. "
176                  f"Valid actions are: {sorted([a.value for a in valid_actions])}"
177              )
178  
179      @property
180      def entity(self) -> WebhookEntity:
181          return self._entity
182  
183      @property
184      def action(self) -> WebhookAction:
185          return self._action
186  
187      @staticmethod
188      def _is_valid_combination(entity: WebhookEntity, action: WebhookAction) -> bool:
189          """
190          Check if an entity/action combination is valid.
191  
192          Args:
193              entity: The webhook entity
194              action: The webhook action
195  
196          Returns:
197              True if the combination is valid, False otherwise
198          """
199          valid_actions = VALID_ENTITY_ACTIONS.get(entity, set())
200          return action in valid_actions
201  
202      @classmethod
203      def from_proto(cls, proto: ProtoWebhookEvent) -> Self:
204          return cls(
205              entity=WebhookEntity.from_proto(proto.entity),
206              action=WebhookAction.from_proto(proto.action),
207          )
208  
209      @classmethod
210      def from_str(cls, event_str: WebhookEventStr) -> Self:
211          """
212          Create a WebhookEvent from a dot-separated string representation.
213  
214          Args:
215              event_str: Valid webhook event string (e.g., "registered_model.created")
216  
217          Returns:
218              A WebhookEvent instance
219          """
220          match event_str.split("."):
221              case [entity_str, action_str]:
222                  try:
223                      entity = WebhookEntity(entity_str)
224                      action = WebhookAction(action_str)
225                      return cls(entity=entity, action=action)
226                  except ValueError as e:
227                      raise MlflowException.invalid_parameter_value(
228                          f"Invalid entity or action in event string: {event_str}. Error: {e}"
229                      )
230              case _:
231                  raise MlflowException.invalid_parameter_value(
232                      f"Invalid event string format: {event_str}. "
233                      "Expected format: 'entity.action' (e.g., 'registered_model.created')"
234                  )
235  
236      def to_proto(self) -> ProtoWebhookEvent:
237          event = ProtoWebhookEvent()
238          event.entity = self.entity.to_proto()
239          event.action = self.action.to_proto()
240          return event
241  
242      def __str__(self) -> str:
243          return f"{self.entity.value}.{self.action.value}"
244  
245      def __eq__(self, other: object) -> bool:
246          if not isinstance(other, WebhookEvent):
247              return False
248          return self.entity == other.entity and self.action == other.action
249  
250      def __hash__(self) -> int:
251          return hash((self.entity, self.action))
252  
253      def __repr__(self) -> str:
254          return f"WebhookEvent(entity={self.entity}, action={self.action})"
255  
256  
257  class Webhook:
258      """
259      MLflow entity for Webhook.
260      """
261  
262      def __init__(
263          self,
264          webhook_id: str,
265          name: str,
266          url: str,
267          events: list[WebhookEvent],
268          creation_timestamp: int,
269          last_updated_timestamp: int,
270          description: str | None = None,
271          status: str | WebhookStatus = WebhookStatus.ACTIVE,
272          secret: str | None = None,
273          workspace: str | None = None,
274      ):
275          """
276          Initialize a Webhook entity.
277  
278          Args:
279              webhook_id: Unique webhook identifier
280              name: Human-readable webhook name
281              url: Webhook endpoint URL
282              events: List of WebhookEvent objects that trigger this webhook
283              creation_timestamp: Creation timestamp in milliseconds since Unix epoch
284              last_updated_timestamp: Last update timestamp in milliseconds since Unix epoch
285              description: Optional webhook description
286              status: Webhook status (ACTIVE or DISABLED)
287              secret: Optional secret key for HMAC signature verification
288              workspace: Workspace the webhook belongs to
289          """
290          super().__init__()
291          self._webhook_id = webhook_id
292          self._name = name
293          self._url = url
294          if not events:
295              raise MlflowException.invalid_parameter_value("Webhook events cannot be empty")
296          self._events = events
297          self._description = description
298          self._status = WebhookStatus(status) if isinstance(status, str) else status
299          self._secret = secret
300          self._creation_timestamp = creation_timestamp
301          self._last_updated_timestamp = last_updated_timestamp
302          self._workspace = resolve_entity_workspace_name(workspace)
303  
304      @property
305      def webhook_id(self) -> str:
306          return self._webhook_id
307  
308      @property
309      def name(self) -> str:
310          return self._name
311  
312      @property
313      def url(self) -> str:
314          return self._url
315  
316      @property
317      def events(self) -> list[WebhookEvent]:
318          return self._events
319  
320      @property
321      def description(self) -> str | None:
322          return self._description
323  
324      @property
325      def status(self) -> WebhookStatus:
326          return self._status
327  
328      @property
329      def secret(self) -> str | None:
330          return self._secret
331  
332      @property
333      def creation_timestamp(self) -> int:
334          return self._creation_timestamp
335  
336      @property
337      def last_updated_timestamp(self) -> int:
338          return self._last_updated_timestamp
339  
340      @property
341      def workspace(self) -> str:
342          return self._workspace
343  
344      @classmethod
345      def from_proto(cls, proto: ProtoWebhook) -> Self:
346          return cls(
347              webhook_id=proto.webhook_id,
348              name=proto.name,
349              url=proto.url,
350              events=[WebhookEvent.from_proto(e) for e in proto.events],
351              description=proto.description or None,
352              status=WebhookStatus.from_proto(proto.status),
353              creation_timestamp=proto.creation_timestamp,
354              last_updated_timestamp=proto.last_updated_timestamp,
355          )
356  
357      def to_proto(self):
358          webhook = ProtoWebhook()
359          webhook.webhook_id = self.webhook_id
360          webhook.name = self.name
361          webhook.url = self.url
362          webhook.events.extend([event.to_proto() for event in self.events])
363          if self.description:
364              webhook.description = self.description
365          webhook.status = self.status.to_proto()
366          webhook.creation_timestamp = self.creation_timestamp
367          webhook.last_updated_timestamp = self.last_updated_timestamp
368          return webhook
369  
370      def __repr__(self) -> str:
371          return (
372              f"Webhook("
373              f"webhook_id='{self.webhook_id}', "
374              f"name='{self.name}', "
375              f"url='{self.url}', "
376              f"status='{self.status}', "
377              f"workspace='{self.workspace}', "
378              f"events={self.events}, "
379              f"creation_timestamp={self.creation_timestamp}, "
380              f"last_updated_timestamp={self.last_updated_timestamp}"
381              f")"
382          )
383  
384  
385  class WebhookTestResult:
386      """
387      MLflow entity for WebhookTestResult.
388      """
389  
390      def __init__(
391          self,
392          success: bool,
393          response_status: int | None = None,
394          response_body: str | None = None,
395          error_message: str | None = None,
396      ):
397          """
398          Initialize a WebhookTestResult entity.
399  
400          Args:
401              success: Whether the test succeeded
402              response_status: HTTP response status code if available
403              response_body: Response body if available
404              error_message: Error message if test failed
405          """
406          self._success = success
407          self._response_status = response_status
408          self._response_body = response_body
409          self._error_message = error_message
410  
411      @property
412      def success(self) -> bool:
413          return self._success
414  
415      @property
416      def response_status(self) -> int | None:
417          return self._response_status
418  
419      @property
420      def response_body(self) -> str | None:
421          return self._response_body
422  
423      @property
424      def error_message(self) -> str | None:
425          return self._error_message
426  
427      @classmethod
428      def from_proto(cls, proto: ProtoWebhookTestResult) -> Self:
429          return cls(
430              success=proto.success,
431              response_status=proto.response_status or None,
432              response_body=proto.response_body or None,
433              error_message=proto.error_message or None,
434          )
435  
436      def to_proto(self) -> ProtoWebhookTestResult:
437          return ProtoWebhookTestResult(
438              success=self.success,
439              response_status=self.response_status,
440              response_body=self.response_body,
441              error_message=self.error_message,
442          )
443  
444      def __repr__(self) -> str:
445          return (
446              f"WebhookTestResult("
447              f"success={self.success!r}, "
448              f"response_status={self.response_status!r}, "
449              f"response_body={self.response_body!r}, "
450              f"error_message={self.error_message!r}"
451              f")"
452          )