diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py
index b3473c6..5478103 100644
--- a/laborious/utils/repository/model_repository.py
+++ b/laborious/utils/repository/model_repository.py
@@ -22,6 +22,250 @@ class MLFlowRepository():
         self.model_serving = ModelServing(tracking_uri=host,
                                           username=username, password=password)
 
+<<<<<<< HEAD
+=======
+    def get_current_data_df(self, current_data: pd.DataFrame, model_name: str, target: str):
+        """
+        Get the current data as a DataFrame and update the prediction and target columns
+
+        Parameters:
+        current_data (pd.DataFrame): the current data
+        model_name (str): the name of the model
+        target (str): the target column
+
+        Returns:
+        DataFrame: the current data as a DataFrame
+
+
+        """
+        predictions = current_data['prediction']
+
+        target = current_data[target]
+        current_data = self.model_serving.get_transformed_data(
+            model_name, current_data, by='model')
+        current_data['prediction'] = predictions
+        current_data['target'] = target
+
+        return pd.DataFrame(current_data).dropna()
+
+    def get_artifact(self, destination: str, search_by: str, run_id: str = None,
+                     model_name: str = None, artifact_name: str = None) -> None:
+        """
+        Get an artifact in MLflow by experiment or model and save it to a
+        destination path using API.
+        If the artifact is searched by model, the latest production version will be used.
+
+        Args:
+            destination: The destination path to save the artifact.
+            search_by: The way to search for the artifact ('experiment' or 'model').
+            run_id: The run ID of the experiment (if search_by is "experiment").
+            model_name: The name of the model (if search_by is "model").
+            artifact_name: The path of the artifact to download.
+
+        Returns:
+            artifact: The artifact(.csv) downloaded from MLflow.
+        """
+
+        self.model_serving.get_artifact(destination=destination, search_by=search_by,
+                                        run_id=run_id, model_name=model_name, artifact_name=artifact_name)
+
+    def calculate_model_metrics(self, real_data, predictions, flag):
+        """
+        Function to calculate the metrics of a model using API
+
+        Parameters:
+        real_data (array): the real data
+        predictions (array): the predictions
+
+        Returns:
+        dict: the metrics of the model including MSE and R2
+        """
+        return self.model_serving.get_model_metrics(
+            reference_data=None, real_data=real_data,
+            predictions=predictions, type_flag=flag)
+
+    def get_experiment_by_run_id(self, run_id: str) -> dict:
+        # Get the run information using the run_id
+        run = mlflow.get_run(run_id)
+
+        # Extract the experiment ID from the run
+        experiment_id = run.info.experiment_id
+
+        # Get the experiment details using the experiment ID
+        experiment = mlflow.get_experiment(experiment_id)
+        experiment_name = experiment.name
+        return experiment_name
+
+    def get_next_run_name(self, model_name: str) -> str:
+        """ 
+        Function to get the next run number of a specific model
+
+        Parameters:
+        model_name (str): the name of the model
+
+        Returns:
+        str: the next run number
+        """
+
+        runs = mlflow.search_runs(
+            experiment_names=[model_name], order_by=["start_time desc"])
+        next_run_number = len(runs) + 1
+        return f"{model_name}-{next_run_number}"
+
+    def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
+        """
+        Retrain a model with new data.
+
+        Parameters:
+            - data (pandas.DataFrame): The new data to use for retraining.
+            - model_name (str): The name of the model to retrain.
+            - metrics_list (list): The metrics to be used to compare the models.
+            - compare_metrics (bool): If True, the retrain will only be considered if
+                the new model is better than the current one.
+                If False, the retrain will always be considered.
+            - split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets.
+                If False, the data will be used as a unique block for retraining.
+            - update_report (bool): If True, a report will be created with the data of the retrained model.
+            - update_transformation (bool): If True, the model will be updated in the MLflow tracking server.
+            - update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server.
+            - shuffle_data (bool): If True, the data will be shuffled before splitting.
+            - model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'.
+
+
+        Returns:
+            - mlflow.sklearn.Model: The retrained prediction model.
+            - mlflow.sklearn.Model: The retrained data model.
+            - mse (float): The mean squared error of the retrained model.
+            - r2 (float): The R-squared score of the retrained model.
+        """
+
+        # load predictor model
+        predictor_uri = f"models:/{model_name}/production"
+        # load transform model
+        latest_production_id = self.model_serving.get_model_run_id(
+            model_name, stage="Production"
+        )
+        transform_uri = self.model_serving.get_model_uri(
+            latest_production_id, prediction=False
+        )
+        # load
+        data_model = mlflow.sklearn.load_model(transform_uri)
+        prediction_model = mlflow.sklearn.load_model(predictor_uri)
+        data_model = data_model.fit(data)
+        treated_data = data_model.predict(data)
+        # align target column with treated_data
+        target_name = data_model.target_variable
+        y = data[target_name]
+        treated_data = pd.merge(
+            treated_data, y, left_index=True, right_index=True)
+        prediction_model = prediction_model.fit(treated_data)
+        # Example usage
+        experiment = self.get_experiment_by_run_id(latest_production_id)
+        pred_model_atributes = vars(prediction_model)  # load class attributes
+        data_model_atributes = vars(data_model)  # load class attributes
+        mlflow.set_experiment(experiment)
+        experiment_description = "Retrain model {model_name} with new data"
+        current_run_name = self.get_next_run_name(experiment)
+        with mlflow.start_run(
+            run_name=current_run_name, description=experiment_description
+        ) as _run:
+            # update transfomation model
+            # fixed parameters
+            for name_atribute, val_atribute in pred_model_atributes.items():
+                if name_atribute != "model":
+                    mlflow.log_param(name_atribute, val_atribute)
+            # update prediction model
+            for name_atribute, val_atribute in data_model_atributes.items():
+                if name_atribute != "model":
+                    mlflow.log_param(name_atribute, val_atribute)
+            # dynamic parameters, including model itself
+            mlflow.sklearn.log_model(data_model, "data_model")
+            file_path = f"laborious/data/raw_data_{model_name}.csv"
+            data.to_csv(
+                f"laborious/data/raw_data_{model_name}.csv", index=True)
+            # log the data raw
+            mlflow.log_artifact(file_path)
+
+            # dynamic parameters, including model itself
+            mlflow.sklearn.log_model(prediction_model, "prediction_model")
+            mlflow.log_param("retrain", True)
+
+            return "Model retrained successfully", experiment
+
+    def get_experiment(self, experiment_name: str) -> int:
+        experiment = mlflow.get_experiment_by_name(experiment_name)
+
+        if experiment is None:
+            raise ValueError(f'Experiment {experiment_name} not found')
+
+        return int(experiment.experiment_id)
+
+    def get_experiment_last_run(self, experiment_id: int) -> str:
+        runs = mlflow.search_runs(
+            experiment_ids=[experiment_id],
+            filter_string="",  # Sem filtro no MLflow ainda
+            output_format="pandas"
+        )
+
+        # Filtrar apenas as runs onde params.retrain == True
+        filtered_runs = runs[runs["params.retrain"] == 'True']
+
+        # Converter a coluna 'end_time' para datetime
+        filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
+
+        # Ordenar o DataFrame de forma descendente pela coluna 'end_time'
+        filtered_runs = filtered_runs.sort_values(
+            by='end_time', ascending=False)
+
+        # Pegar a última run_id do DataFrame filtrado e ordenado
+        latest_run_id = filtered_runs.iloc[0]['run_id']
+
+        return latest_run_id
+
+    def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
+        # Registrar o modelo
+        # Aqui estamos assumindo que você já tem um modelo salvo,
+        # caso contrário você precisará treiná-lo e salvá-lo primeiro.
+        # Se o modelo já está registrado, você pode usar o método
+        # register_model() ou pyfunc.load_model() para isso.
+        mlflow.register_model(
+            f"runs:/{run_id}/prediction_model", model_name)
+
+        # Colocar a versão do modelo em produção
+        # Depois de registrar o modelo, precisamos pegar a versão mais
+        # recente do modelo e movê-lo para o estágio 'Production'
+        client = mlflow.tracking.MlflowClient()
+
+        # Obter a versão mais recente registrada do modelo
+        model_versions = client.get_registered_model(
+            model_name).latest_versions
+        max_version = max(model_versions, key=lambda x: int(x.version)).version
+
+        # Mover a versão mais recente do modelo para o estágio de 'Production'
+        client.transition_model_version_stage(
+            name=model_name,
+            version=max_version,
+            stage="Production",
+            archive_existing_versions=True
+        )
+
+        return {
+            'model_name': model_name,
+            'version': max_version,
+            'mlflow_run_id': run_id
+        }
+
+    def update_production_model(self, experiment: str, model_name: str) -> dict:
+
+        experiment_id = self.get_experiment(experiment)
+        run_id = self.get_experiment_last_run(experiment_id)
+        metadata = self.update_production_model_by_run_id(run_id, model_name)
+
+        metadata['mlflow_experiment_id'] = experiment_id
+
+        return metadata
+
+>>>>>>> main
     def transform(self, model_name: str, data: pd.DataFrame, model_retention: int):
         """
         Transform data using a model.