from image_prediction.model_loader.database.connector import DatabaseConnector
from image_prediction.redai_adapter.mlflow import MlflowModelReader


class MlflowConnector(DatabaseConnector):
    def __init__(self, mlflow_reader: MlflowModelReader):
        self.mlflow_reader = mlflow_reader

    def get_object(self, run_id):
        return self.mlflow_reader[run_id]