From 0916b5dce0f3990d10740096df0116157b93a10a Mon Sep 17 00:00:00 2001
From: shekhar-rudder <85345786+shekhar-rudder@users.noreply.github.com>
Date: Thu, 11 Jul 2024 01:01:36 +0530
Subject: [PATCH] fix: move features field inside model_spec for prediction
 model (#369)

---
 samples/py_native/models/profiles.yaml               | 12 ++++++------
 .../profiles_mlcorelib/py_native/prediction.py       |  6 +++++-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/samples/py_native/models/profiles.yaml b/samples/py_native/models/profiles.yaml
index 0d5f35410..5d1434913 100644
--- a/samples/py_native/models/profiles.yaml
+++ b/samples/py_native/models/profiles.yaml
@@ -43,9 +43,9 @@ models:
           column_names:
             percentile: percentile_classification_churn_7_days
             score: classification_churn_7_days
-    features:
-      - name: percentile_classification_churn_7_days
-        description: description
+      features:
+        - name: percentile_classification_churn_7_days
+          description: description
   - name: prediction_regression_model
     model_type: prediction_model
     model_spec:
@@ -65,6 +65,6 @@ models:
           column_names:
             percentile: percentile_regression_days_since_last_seen
             score: regression_days_since_last_seen
-    features:
-      - name: percentile_regression_days_since_last_seen
-        description: description
+      features:
+        - name: percentile_regression_days_since_last_seen
+          description: description
diff --git a/src/predictions/profiles_mlcorelib/py_native/prediction.py b/src/predictions/profiles_mlcorelib/py_native/prediction.py
index 630fa4d57..f2ab9aed0 100644
--- a/src/predictions/profiles_mlcorelib/py_native/prediction.py
+++ b/src/predictions/profiles_mlcorelib/py_native/prediction.py
@@ -3,7 +3,10 @@
 from profiles_rudderstack.material import WhtMaterial
 from profiles_rudderstack.logger import Logger
 from typing import Tuple
-from profiles_rudderstack.schema import EntityKeyBuildSpecSchema
+from profiles_rudderstack.schema import (
+    EntityKeyBuildSpecSchema,
+    FeatureDetailsBuildSpecSchema,
+)
 
 from .warehouse import standardize_ref_name
 
@@ -24,6 +27,7 @@ class PredictionModel(BaseModelType):
         "properties": {
             "occurred_at_col": {"type": "string"},
             **EntityKeyBuildSpecSchema["properties"],
+            **FeatureDetailsBuildSpecSchema["properties"],
             "validity_time": {"type": "string"},
             "inputs": {"type": "array", "items": {"type": "string"}, "minItems": 1},
             "training_model": {"type": "string"},