chore: fix mypy errors in task-plugins/dioptra_builtins

task-plugins/dioptra_builtins/artifacts/mlflow.py

@@ -25,6 +25,7 @@
 from typing import Any, Callable, Dict, Optional, Union
 
 import mlflow
+import mlflow.log_artifact
 import pandas as pd
 import structlog
 from mlflow.tracking import MlflowClient

task-plugins/dioptra_builtins/registry/mlflow.py

@@ -21,6 +21,7 @@
 from typing import Optional
 
 import mlflow
+import mlflow.keras
 import structlog
 from mlflow.entities import Run as MlflowRun
 from mlflow.entities.model_registry import ModelVersion

task-plugins/dioptra_builtins/tracking/mlflow.py

@@ -21,6 +21,8 @@
 from typing import Dict
 
 import mlflow
+import mlflow.log_metric
+import mlflow.log_param
 import structlog
 from structlog.stdlib import BoundLogger