fix: move features field inside model_spec for prediction model (#369)

rudderlabs · Jul 10, 2024 · 0916b5d · 0916b5d
1 parent 4782cd9
commit 0916b5d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git 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
@@ -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"},