Fix null cases for openai prompt and embeddings (#2457)

Author: ranadeepsingh

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIEmbedding.scala

@@ -16,7 +16,7 @@ import org.apache.spark.ml.functions.array_to_vector
 import org.apache.spark.sql.{DataFrame, Dataset, Row}
 import org.apache.spark.sql.types._
 import scala.language.existentials
-import org.apache.spark.sql.functions.{col, element_at, struct}
+import org.apache.spark.sql.functions.{col, element_at, struct, when}
 import spray.json.DefaultJsonProtocol._
 import spray.json._
 import HasReturnUsage.UsageMappings
@@ -101,13 +101,18 @@ class OpenAIEmbedding (override val uid: String) extends OpenAIServicesBase(uid)
       )
       parsed.withColumn(
         getOutputCol,
-        struct(
-          vectorCol.alias("response"),
-          usageCol.alias("usage")
+        when(responseCol.isNotNull,
+          struct(
+            vectorCol.alias("response"),
+            usageCol.alias("usage")
+          )
         )
       )
     } else {
-      parsed.withColumn(getOutputCol, vectorCol)
+      parsed.withColumn(
+        getOutputCol,
+        when(responseCol.isNotNull, vectorCol)
+      )
     }
   }
 

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala

@@ -277,10 +277,15 @@ class OpenAIPrompt(override val uid: String) extends Transformer
         val usageCol = normalizeUsageColumn(responseCol.getField("usage"), mapping)
         df.withColumn(
           getOutputCol,
-          F.struct(parsedCol.alias("response"), usageCol.alias("usage"))
+          F.when(parsedCol.isNotNull,
+            F.struct(parsedCol.alias("response"), usageCol.alias("usage"))
+          )
         )
       case None =>
-        df.withColumn(getOutputCol, parsedCol)
+        df.withColumn(
+          getOutputCol,
+          F.when(parsedCol.isNotNull, parsedCol)
+        )
     }
   }
 

cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIEmbeddingsSuite.scala

@@ -173,6 +173,36 @@ class OpenAIEmbeddingsSuite extends TransformerFuzzing[OpenAIEmbedding] with Ope
     assert(outputDetails != null)
   }
 
+  test("null input returns null output with returnUsage false") {
+    val dfWithNull = Seq(
+      Some("Once upon a time"),
+      None,
+      Some("SynapseML is ")
+    ).toDF("text")
+
+    val e = usageEmbedding("null_test_vec")
+    val results = e.transform(dfWithNull).collect()
+
+    assert(results(0).getAs[Vector]("null_test_vec") != null)
+    assert(results(1).getAs[Vector]("null_test_vec") == null)
+    assert(results(2).getAs[Vector]("null_test_vec") != null)
+  }
+
+  test("null input returns null output with returnUsage true") {
+    val dfWithNull = Seq(
+      Some("Once upon a time"),
+      None,
+      Some("SynapseML is ")
+    ).toDF("text")
+
+    val e = usageEmbedding("null_test_usage").setReturnUsage(true)
+    val results = e.transform(dfWithNull).collect()
+
+    assert(results(0).getAs[Row]("null_test_usage") != null)
+    assert(results(1).getAs[Row]("null_test_usage") == null)
+    assert(results(2).getAs[Row]("null_test_usage") != null)
+  }
+
 
   override def testObjects(): Seq[TestObject[OpenAIEmbedding]] =
     Seq(new TestObject(embedding, df))

cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPromptSuite.scala

@@ -204,42 +204,38 @@ class OpenAIPromptSuite extends TransformerFuzzing[OpenAIPrompt] with OpenAIAPIK
     assert(outputDetails != null)
   }
 
-  lazy val promptGpt4: OpenAIPrompt = new OpenAIPrompt()
-    .setSubscriptionKey(openAIAPIKey)
-    .setDeploymentName(deploymentName)
-    .setCustomServiceName(openAIServiceName)
-    .setOutputCol("outParsed")
-    .setTemperature(0)
+  test("null input returns null output with returnUsage false") {
+    val dfWithNull = Seq(
+      (Some("apple"), "fruits"),
+      (None, "cars"),
+      (Some("cake"), "dishes")
+    ).toDF("text", "category")
+
+    val p = usagePrompt("null_test")
+    val results = p.transform(dfWithNull).select("null_test").collect()
+
+    assert(results(0).getSeq[String](0) != null)
+    assert(results(1).get(0) == null)
+    assert(results(2).getSeq[String](0) != null)
+  }
 
-  test("Basic Usage - Gpt 4") {
-    val nonNullCount = promptGpt4
-      .setPromptTemplate("give me a comma separated list of 5 {category}, starting with {text} ")
-      .setPostProcessing("csv")
-      .transform(df)
-      .select("outParsed")
-      .collect()
-      .count(r => Option(r.getSeq[String](0)).isDefined)
+  test("null input returns null output with returnUsage true") {
+    val dfWithNull = Seq(
+      (Some("apple"), "fruits"),
+      (None, "cars"),
+      (Some("cake"), "dishes")
+    ).toDF("text", "category")
 
-    assert(nonNullCount == 3)
-  }
+    val p = usagePrompt("null_test_usage").setReturnUsage(true)
+    val results = p.transform(dfWithNull).select("null_test_usage").collect()
 
-  test("Basic Usage JSON - Gpt 4") {
-    promptGpt4.setPromptTemplate(
-        """Split a word into prefix and postfix a respond in JSON
-          |Cherry: {{"prefix": "Che", "suffix": "rry"}}
-          |{text}:
-          |""".stripMargin)
-      .setPostProcessing("json")
-      .setPostProcessingOptions(Map("jsonSchema" -> "prefix STRING, suffix STRING"))
-      .transform(df)
-      .select("outParsed")
-      .where(col("outParsed").isNotNull)
-      .collect()
-      .foreach(r => assert(r.getStruct(0).getString(0).nonEmpty))
+    assert(results(0).getStruct(0) != null)
+    assert(results(1).get(0) == null)
+    assert(results(2).getStruct(0) != null)
   }
 
-  test("Basic Usage JSON - Gpt 4 without explicit post-processing") {
-    promptGpt4.setPromptTemplate(
+  test("Basic Usage JSON - without explicit post-processing") {
+    prompt.setPromptTemplate(
                 """Split a word into prefix and postfix a respond in JSON
                   |Cherry: {{"prefix": "Che", "suffix": "rry"}}
                   |{text}:
@@ -252,8 +248,8 @@ class OpenAIPromptSuite extends TransformerFuzzing[OpenAIPrompt] with OpenAIAPIK
               .foreach(r => assert(r.getStruct(0).getString(0).nonEmpty))
   }
 
-  test("Setting and Keeping Messages Col - Gpt 4") {
-    promptGpt4.setMessagesCol("messages")
+  test("Setting and Keeping Messages Col") {
+    prompt.setMessagesCol("messages")
       .setDropPrompt(false)
       .setPromptTemplate(
         """Classify each word as to whether they are an F1 team or not
@@ -268,23 +264,23 @@ class OpenAIPromptSuite extends TransformerFuzzing[OpenAIPrompt] with OpenAIAPIK
       .foreach(r => assert(r.get(0) != null))
   }
 
-  test("Basic Usage JSON - Gpt 4o with responseFormat") {
-    val promptGpt4o: OpenAIPrompt = new OpenAIPrompt()
+  test("json_object Response Format Usage") {
+    val promptJSONObject: OpenAIPrompt = new OpenAIPrompt()
       .setSubscriptionKey(openAIAPIKey)
       .setDeploymentName(deploymentName)
       .setCustomServiceName(openAIServiceName)
       .setOutputCol("outParsed")
       .setTemperature(0)
       .setPromptTemplate(
-        """Split a word into prefix and postfix
+        """Split a word into prefix and postfix in JSON format
           |Cherry: {{"prefix": "Che", "suffix": "rry"}}
           |{text}:
           |""".stripMargin)
       .setResponseFormat("json_object")
       .setPostProcessingOptions(Map("jsonSchema" -> "prefix STRING, suffix STRING"))
 
 
-    promptGpt4o.transform(df)
+    promptJSONObject.transform(df)
                .select("outParsed")
                .where(col("outParsed").isNotNull)
                .collect()
@@ -331,21 +327,21 @@ class OpenAIPromptSuite extends TransformerFuzzing[OpenAIPrompt] with OpenAIAPIK
     lazy val customRootUrlValue: String = sys.env.getOrElse("CUSTOM_ROOT_URL", "")
     lazy val customHeadersValues: Map[String, String] = Map("X-ModelType" -> "gpt-4-turbo-chat-completions")
 
-    lazy val customPromptGpt4: OpenAIPrompt = new OpenAIPrompt()
+    lazy val customPrompt: OpenAIPrompt = new OpenAIPrompt()
       .setCustomUrlRoot(customRootUrlValue)
       .setOutputCol("outParsed")
       .setTemperature(0)
 
     if (accessToken.isEmpty) {
-      customPromptGpt4.setSubscriptionKey(openAIAPIKey)
+      customPrompt.setSubscriptionKey(openAIAPIKey)
         .setDeploymentName(deploymentName)
         .setCustomServiceName(openAIServiceName)
     } else {
-      customPromptGpt4.setAADToken(accessToken)
+      customPrompt.setAADToken(accessToken)
         .setCustomHeaders(customHeadersValues)
     }
 
-    customPromptGpt4.setPromptTemplate("give me a comma separated list of 5 {category}, starting with {text} ")
+    customPrompt.setPromptTemplate("give me a comma separated list of 5 {category}, starting with {text} ")
       .setPostProcessing("csv")
       .transform(df)
       .select("outParsed")