postgresml · SilasMarvin · Dec 1, 2023 · Nov 22, 2023 · Nov 22, 2023 · Nov 28, 2023
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -632,6 +632,50 @@ pub fn transform_string(
     }
 }
 
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "transform")]
+#[allow(unused_variables)] // cache is maintained for api compatibility
+pub fn transform_conversational_json(
+    task: JsonB,
+    args: default!(JsonB, "'{}'"),
+    inputs: default!(Vec<JsonB>, "ARRAY[]::JSONB[]"),
+    cache: default!(bool, false),
+) -> JsonB {
+    if !task.0["task"]
+        .as_str()
+        .is_some_and(|v| v == "conversational")
+    {
+        error!(
+            "ARRAY[]::JSONB inputs for transformer should only be used with a conversational task"
+        );
+    }
+    match crate::bindings::transformers::transform(&task.0, &args.0, inputs) {
+        Ok(output) => JsonB(output),
+        Err(e) => error!("{e}"),
+    }
+}
+
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "transform")]
+#[allow(unused_variables)] // cache is maintained for api compatibility
+pub fn transform_conversational_string(
+    task: String,
+    args: default!(JsonB, "'{}'"),
+    inputs: default!(Vec<JsonB>, "ARRAY[]::JSONB[]"),
+    cache: default!(bool, false),
+) -> JsonB {
+    if task != "conversational" {
+        error!(
+            "ARRAY[]::JSONB inputs for transformer should only be used with a conversational task"
+        );
+    }
+    let task_json = json!({ "task": task });
+    match crate::bindings::transformers::transform(&task_json, &args.0, inputs) {
+        Ok(output) => JsonB(output),
+        Err(e) => error!("{e}"),
+    }
+}
+
 #[cfg(all(feature = "python", not(feature = "use_as_lib")))]
 #[pg_extern(immutable, parallel_safe, name = "transform_stream")]
 #[allow(unused_variables)] // cache is maintained for api compatibility
@@ -640,7 +684,7 @@ pub fn transform_stream_json(
     args: default!(JsonB, "'{}'"),
     input: default!(&str, "''"),
     cache: default!(bool, false),
-) -> SetOfIterator<'static, String> {
+) -> SetOfIterator<'static, JsonB> {
     // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
     let python_iter =
         crate::bindings::transformers::transform_stream_iterator(&task.0, &args.0, input)
@@ -657,7 +701,7 @@ pub fn transform_stream_string(
     args: default!(JsonB, "'{}'"),
     input: default!(&str, "''"),
     cache: default!(bool, false),
-) -> SetOfIterator<'static, String> {
+) -> SetOfIterator<'static, JsonB> {
     let task_json = json!({ "task": task });
     // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
     let python_iter =
@@ -667,6 +711,54 @@ pub fn transform_stream_string(
     SetOfIterator::new(python_iter)
 }
 
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "transform_stream")]
+#[allow(unused_variables)] // cache is maintained for api compatibility
+pub fn transform_stream_conversational_json(
+    task: JsonB,
+    args: default!(JsonB, "'{}'"),
+    inputs: default!(Vec<JsonB>, "ARRAY[]::JSONB[]"),
+    cache: default!(bool, false),
+) -> SetOfIterator<'static, JsonB> {
+    if !task.0["task"]
+        .as_str()
+        .is_some_and(|v| v == "conversational")
+    {
+        error!(
+            "ARRAY[]::JSONB inputs for transformer_stream should only be used with a conversational task"
+        );
+    }
+    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    let python_iter =
+        crate::bindings::transformers::transform_stream_iterator(&task.0, &args.0, inputs)
+            .map_err(|e| error!("{e}"))
+            .unwrap();
+    SetOfIterator::new(python_iter)
+}
+
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "transform_stream")]
+#[allow(unused_variables)] // cache is maintained for api compatibility
+pub fn transform_stream_conversational_string(
+    task: String,
+    args: default!(JsonB, "'{}'"),
+    inputs: default!(Vec<JsonB>, "ARRAY[]::JSONB[]"),
+    cache: default!(bool, false),
+) -> SetOfIterator<'static, JsonB> {
+    if task != "conversational" {
+        error!(
+            "ARRAY::JSONB inputs for transformer_stream should only be used with a conversational task"
+        );
+    }
+    let task_json = json!({ "task": task });
+    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    let python_iter =
+        crate::bindings::transformers::transform_stream_iterator(&task_json, &args.0, inputs)
+            .map_err(|e| error!("{e}"))
+            .unwrap();
+    SetOfIterator::new(python_iter)
+}
+
 #[cfg(feature = "python")]
 #[pg_extern(immutable, parallel_safe, name = "generate")]
 fn generate(project_name: &str, inputs: &str, config: default!(JsonB, "'{}'")) -> String {

diff --git a/pgml-extension/src/bindings/transformers/transform.rs b/pgml-extension/src/bindings/transformers/transform.rs
@@ -23,28 +23,30 @@ impl TransformStreamIterator {
 }
 
 impl Iterator for TransformStreamIterator {
-    type Item = String;
+    type Item = JsonB;
     fn next(&mut self) -> Option<Self::Item> {
         // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
-        Python::with_gil(|py| -> Result<Option<String>, PyErr> {
+        Python::with_gil(|py| -> Result<Option<JsonB>, PyErr> {
             let code = "next(python_iter)";
             let res: &PyAny = py.eval(code, Some(self.locals.as_ref(py)), None)?;
             if res.is_none() {
                 Ok(None)
             } else {
-                let res: String = res.extract()?;
-                Ok(Some(res))
+                eprintln!("\nHERE WE ARE!\n");
+                let res: Vec<String> = res.extract()?;
+                eprintln!("\nYUP WE DIDNT GET HERE\n");
+                Ok(Some(JsonB(serde_json::to_value(res).unwrap())))
             }
         })
         .map_err(|e| error!("{e}"))
         .unwrap()
     }
 }
 
-pub fn transform(
+pub fn transform<T: serde::Serialize>(
     task: &serde_json::Value,
     args: &serde_json::Value,
-    inputs: Vec<&str>,
+    inputs: T,
 ) -> Result<serde_json::Value> {
     crate::bindings::python::activate()?;
     whitelist::verify_task(task)?;
@@ -74,17 +76,17 @@ pub fn transform(
     Ok(serde_json::from_str(&results)?)
 }
 
-pub fn transform_stream(
+pub fn transform_stream<T: serde::Serialize>(
     task: &serde_json::Value,
     args: &serde_json::Value,
-    input: &str,
+    input: T,
 ) -> Result<Py<PyAny>> {
     crate::bindings::python::activate()?;
     whitelist::verify_task(task)?;
 
     let task = serde_json::to_string(task)?;
     let args = serde_json::to_string(args)?;
-    let inputs = serde_json::to_string(&vec![input])?;
+    let input = serde_json::to_string(&input)?;
 
     Python::with_gil(|py| -> Result<Py<PyAny>> {
         let transform: Py<PyAny> = get_module!(PY_MODULE)
@@ -99,7 +101,7 @@ pub fn transform_stream(
                     &[
                         task.into_py(py),
                         args.into_py(py),
-                        inputs.into_py(py),
+                        input.into_py(py),
                         true.into_py(py),
                     ],
                 ),
@@ -110,10 +112,10 @@ pub fn transform_stream(
     })
 }
 
-pub fn transform_stream_iterator(
+pub fn transform_stream_iterator<T: serde::Serialize>(
     task: &serde_json::Value,
     args: &serde_json::Value,
-    input: &str,
+    input: T,
 ) -> Result<TransformStreamIterator> {
     let python_iter = transform_stream(task, args, input)
         .map_err(|e| error!("{e}"))