Move task verification out of internal binding (#1368)

levkk · web-flow · commit 0307343e96cb · 2024-03-12T19:14:34.000-07:00
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -647,6 +647,10 @@ pub fn transform_json(
     inputs: default!(Vec<&str>, "ARRAY[]::TEXT[]"),
     cache: default!(bool, false),
 ) -> JsonB {
+    if let Err(err) = crate::bindings::transformers::whitelist::verify_task(&task.0) {
+        error!("{err}");
+    }
+
     match crate::bindings::transformers::transform(&task.0, &args.0, inputs) {
         Ok(output) => JsonB(output),
         Err(e) => error!("{e}"),
@@ -663,6 +667,9 @@ pub fn transform_string(
     cache: default!(bool, false),
 ) -> JsonB {
     let task_json = json!({ "task": task });
+    if let Err(err) = crate::bindings::transformers::whitelist::verify_task(&task_json) {
+        error!("{err}");
+    }
     match crate::bindings::transformers::transform(&task_json, &args.0, inputs) {
         Ok(output) => JsonB(output),
         Err(e) => error!("{e}"),
@@ -681,6 +688,9 @@ pub fn transform_conversational_json(
     if !task.0["task"].as_str().is_some_and(|v| v == "conversational") {
         error!("ARRAY[]::JSONB inputs for transform should only be used with a conversational task");
     }
+    if let Err(err) = crate::bindings::transformers::whitelist::verify_task(&task.0) {
+        error!("{err}");
+    }
     match crate::bindings::transformers::transform(&task.0, &args.0, inputs) {
         Ok(output) => JsonB(output),
         Err(e) => error!("{e}"),
@@ -700,6 +710,9 @@ pub fn transform_conversational_string(
         error!("ARRAY[]::JSONB inputs for transform should only be used with a conversational task");
     }
     let task_json = json!({ "task": task });
+    if let Err(err) = crate::bindings::transformers::whitelist::verify_task(&task_json) {
+        error!("{err}");
+    }
     match crate::bindings::transformers::transform(&task_json, &args.0, inputs) {
         Ok(output) => JsonB(output),
         Err(e) => error!("{e}"),
diff --git a/pgml-extension/src/bindings/transformers/transform.rs b/pgml-extension/src/bindings/transformers/transform.rs
@@ -46,8 +46,6 @@ pub fn transform<T: serde::Serialize>(
     args: &serde_json::Value,
     inputs: T,
 ) -> Result<serde_json::Value> {
-    whitelist::verify_task(task)?;
-
     let task = serde_json::to_string(task)?;
     let args = serde_json::to_string(args)?;
     let inputs = serde_json::to_string(&inputs)?;
