rafactor BaseMetricRequest to use Pydantic BaseModel and update usage

Author: saichandrapandraju

src/service/payloads/metrics/base_metric_request.py

@@ -1,23 +1,19 @@
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from typing import Dict, Optional
+from pydantic import BaseModel, ConfigDict
 
 
-class BaseMetricRequest(ABC):
+class BaseMetricRequest(BaseModel):
     """
     Abstract base class for metric requests.
     """
+    # To allow extra fields to be set on instances
+    model_config = ConfigDict(extra='allow')
 
-    def __init__(
-        self,
-        model_id: str,
-        metric_name: str,
-        request_name: Optional[str],
-        batch_size: Optional[int],
-    ) -> None:
-        self.model_id: str = model_id
-        self.metric_name: str = metric_name
-        self.request_name: str = request_name
-        self.batch_size: int = batch_size
+    model_id: str
+    metric_name: str
+    request_name: Optional[str] = None
+    batch_size: Optional[int] = None
 
     def retrieve_default_tags(self) -> Dict[str, str]:
         output: Dict[str, str] = {}

src/service/payloads/metrics/request_reconciler.py

@@ -39,7 +39,12 @@ def reconcile_with_metadata(
             storage_metadata: The storage metadata to use for reconciliation
         """
         try:
-            for name in dir(request.__class__):
+            # Get both model fields and dynamically added instance attributes
+            model_field_names = set(request.__class__.model_fields.keys())
+            instance_attribute_names = set(request.__dict__.keys())
+            all_field_names = model_field_names.union(instance_attribute_names)
+
+            for name in all_field_names:
                 if name.startswith("_"):
                     continue
 

tests/service/payloads/metrics/test_request_reconciler.py

@@ -21,7 +21,12 @@ class MockReconcilableRequest(BaseMetricRequest):
     """Mock request class with reconcilable fields."""
 
     def __init__(self, model_id: str = "test_model"):
-        super().__init__(model_id, "test_metric", "test_request", 100)
+        super().__init__(
+            model_id=model_id,
+            metric_name="test_metric",
+            request_name="test_request",
+            batch_size=100,
+        )
 
         # Create reconcilable feature field
         self.protected_attribute = ReconcilableFeature(