fix definition.py

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_monitoring/definition.py

@@ -20,8 +20,11 @@
 from azure.ai.ml.entities._monitoring.alert_notification import AlertNotification
 from azure.ai.ml.entities._monitoring.compute import ServerlessSparkCompute
 from azure.ai.ml.entities._monitoring.signals import (
+    CustomMonitoringSignal,
     DataDriftSignal,
     DataQualitySignal,
+    FeatureAttributionDriftSignal,
+    GenerationSafetyQualitySignal,
     MonitoringSignal,
     PredictionDriftSignal,
 )
@@ -62,7 +65,17 @@ def __init__(
         *,
         compute: ServerlessSparkCompute,
         monitoring_target: Optional[MonitoringTarget] = None,
-        monitoring_signals: Optional[Dict] = None,
+        monitoring_signals: Dict[
+            str,
+            Union[
+                DataDriftSignal,
+                DataQualitySignal,
+                PredictionDriftSignal,
+                FeatureAttributionDriftSignal,
+                CustomMonitoringSignal,
+                GenerationSafetyQualitySignal,
+            ],
+        ] = None,  # type: ignore[assignment]
         alert_notification: Optional[Union[Literal["azmonitoring"], AlertNotification]] = None,
     ) -> None:
         self.compute = compute
@@ -108,16 +121,19 @@ def _from_rest_object(
                 from_rest_alert_notification = AZMONITORING
             else:
                 from_rest_alert_notification = AlertNotification._from_rest_object(obj.alert_notification_setting)
+
+        _monitoring_signals = {}
+        for signal_name, signal in obj.signals.items():
+            _monitoring_signals[signal_name] = MonitoringSignal._from_rest_object(signal)
+
         return cls(
             compute=ServerlessSparkCompute._from_rest_object(obj.compute_configuration),
             monitoring_target=MonitoringTarget(
                 endpoint_deployment_id=obj.monitoring_target.deployment_id, ml_task=obj.monitoring_target.task_type
             )
             if obj.monitoring_target
             else None,
-            monitoring_signals={
-                signal_name: MonitoringSignal._from_rest_object(signal) for signal_name, signal in obj.signals.items()
-            },
+            monitoring_signals=_monitoring_signals,  # type: ignore[arg-type]
             alert_notification=from_rest_alert_notification,
         )