feat: add output contract in prediction model (#371)

rudderlabs · Jul 11, 2024 · f31f151 · f31f151
1 parent 0916b5d
commit f31f151
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/predictions/profiles_mlcorelib/py_native/prediction.py b/src/predictions/profiles_mlcorelib/py_native/prediction.py
@@ -64,6 +64,17 @@ class PredictionModel(BaseModelType):
     }
 
     def __init__(self, build_spec: dict, schema_version: int, pb_version: str) -> None:
+        build_spec["contract"] = {
+            "with_entity_ids": [build_spec["entity_key"]],
+            "with_columns": [
+                {
+                    "name": build_spec["ml_config"]["outputs"]["column_names"][
+                        "percentile"
+                    ]
+                },
+                {"name": build_spec["ml_config"]["outputs"]["column_names"]["score"]},
+            ],
+        }
         super().__init__(build_spec, schema_version, pb_version)
 
     def get_material_recipe(self) -> PyNativeRecipe: