From 8d4dc65aaa054ab01b5a5d8c403c9fcb500c7612 Mon Sep 17 00:00:00 2001
From: vitor-aignosi <vitor.santos@aignosi.com.br>
Date: Thu, 16 Apr 2026 10:26:42 -0300
Subject: [PATCH] feat: enhance report generation with target name and content
 injection

- Added target_name parameter to Reports class for improved report context.
- Updated inject_content function to ensure proper handling of HTML sections.
- Modified DataManagerRepository to join predictions with training and validation data for accurate report generation.
---
 model_manager/sientia/reports.py              | 26 +++++++++++++++----
 .../repository/data_manager_repository.py     | 24 ++++++++++++++---
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/model_manager/sientia/reports.py b/model_manager/sientia/reports.py
index adb8ffd..0d9b657 100644
--- a/model_manager/sientia/reports.py
+++ b/model_manager/sientia/reports.py
@@ -2,7 +2,7 @@ import os
 from collections.abc import Sequence
 from typing import Any
 
-from bs4 import BeautifulSoup
+from bs4 import BeautifulSoup, Tag
 from evidently.metric_preset import DataDriftPreset
 from evidently.metrics import (
     ColumnSummaryMetric,
@@ -20,6 +20,7 @@ from evidently.metrics import (
 from evidently.metrics.base_metric import generate_column_metrics
 from evidently.options import ColorOptions
 from evidently.report import Report
+from evidently.pipeline.column_mapping import ColumnMapping
 
 COLOR_DISCRETE_SEQUENCE = (
     '#ed0400',
@@ -44,9 +45,14 @@ def load_html_from_file(file_path):
 def inject_content(main_html, section_id, content):
     soup = BeautifulSoup(main_html, 'html.parser')
     section = soup.find(id=section_id)
-    if section:
+    
+    # Verifica se a seção foi encontrada E se ela é uma Tag (não uma string)
+    if section and isinstance(section, Tag):
         section.clear()
-        section.append(BeautifulSoup(content, 'html.parser'))
+        # Converte o conteúdo para um fragmento de BeautifulSoup e anexa
+        new_content = BeautifulSoup(content, 'html.parser')
+        section.append(new_content)
+        
     return str(soup)
 
 
@@ -68,7 +74,7 @@ class Reports:
     """
 
     def __init__(
-        self, reference_data: Any, current_data: Any, base_path: str | None = None
+        self, reference_data: Any, current_data: Any, target_name: str, base_path: str | None = None
     ) -> None:
         """
         Initializes an instance of the AigReport class.
@@ -84,6 +90,7 @@ class Reports:
         self.report: Any = None
         self.ref_data = reference_data
         self.cur_data = current_data
+        self.target_name = target_name
         self.set_color_options(primary_color='#0F4C81', secondary_color='#001E60')
         self.base_path = base_path
 
@@ -145,8 +152,17 @@ class Reports:
         ]
         self.metrics.extend(metrics)
         if run:
+            mapping = ColumnMapping()
+
+            mapping.target = self.target_name
+            mapping.prediction = 'prediction'
+            
             report = Report(metrics=metrics, options=self.options)
-            report.run(reference_data=self.ref_data, current_data=self.cur_data)
+            report.run(
+                reference_data=self.ref_data,
+                current_data=self.cur_data,
+                column_mapping=mapping,
+            )
             self.sections['regression'] = report.as_dict()
             if self.base_path:
                 # Note: Relies on Evidently's save_html() to properly manage file I/O
diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py
index 2f6f387..a87e7f5 100644
--- a/model_manager/utils/repository/data_manager_repository.py
+++ b/model_manager/utils/repository/data_manager_repository.py
@@ -511,11 +511,26 @@ class DataManagerRepository(SientiaMonitoring):
         if data.run_name is None:
             raise ValueError('run_name is not set, cannot generate report')
 
-        # Convert data to float64 for report generation
-        # This may raise ValueError if data contains non-numeric values
-        reference_data = data.train_data
-        current_data = data.val_data
+        if data.y_train_pred is None or data.y_pred is None:
+            raise ValueError('y_train_pred or y_pred is not set, cannot generate report')
+
+
+        y_train_pred = data.y_train_pred.rename(
+            columns={data.params.target_variable: 'prediction'}
+        )
+        y_val_pred = data.y_pred.rename(
+            columns={data.params.target_variable: 'prediction'}
+        )
+        
+        # Join the predictions to the data
+        reference_data = y_train_pred[['prediction']].join(
+            data.train_data, how='inner'
+        )
         reference_data_float = reference_data.astype(np.float64)
+
+        current_data = y_val_pred[['prediction']].join(
+            data.val_data, how='inner'
+        )
         current_data_float = current_data.astype(np.float64)
 
         # Initialize report generator
@@ -524,6 +539,7 @@ class DataManagerRepository(SientiaMonitoring):
             reference_data=reference_data_float,
             current_data=current_data_float,
             base_path=data.run_dir,
+            target_name=data.params.target_variable,
         )
 
         # Generate report sections