OpenAIChatCompletion class method shouldSkip was not working as desig…

…ned as there might OpenAIMessages with empty contents. These rows can be identified before sending requests to OpenAI service and getting errors. For now, we are skipping any rows that do not have any user message

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

@@ -46,8 +46,20 @@ class OpenAIChatCompletion(override val uid: String) extends OpenAIServicesBase(
 
   override val subscriptionKeyHeaderName: String = "api-key"
 
-  override def shouldSkip(row: Row): Boolean =
-    super.shouldSkip(row) || Option(row.getAs[Row](getMessagesCol)).isEmpty
+  /**
+   * Check if the row should be skipped. A row should be skipped if the messages column is empty or if any of the
+   * messages are empty or if there not exists a single message with "role" as "user" and "content" as non-empty.
+   */
+  override def shouldSkip(row: Row): Boolean ={
+    super.shouldSkip(row) || {
+      val messages = Option(row.getAs[Seq[Row]](getMessagesCol)).getOrElse(Seq.empty)
+      messages.isEmpty || !messages.exists(m =>{
+        val content = Option(m.getAs[String]("content"))
+        val role = Option(m.getAs[String]("role"))
+        role.nonEmpty && role.get == "user" && content.nonEmpty && content.get.nonEmpty
+      })
+    }
+  }
 
   override protected def getVectorParamMap: Map[String, String] = super.getVectorParamMap
     .updated("messages", getMessagesCol)

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

@@ -47,6 +47,11 @@ class OpenAIChatCompletionSuite extends TransformerFuzzing[OpenAIChatCompletion]
       OpenAIMessage("system", "You are very excited"),
       OpenAIMessage("user", null) //scalastyle:ignore null
     ),
+    Seq(
+      OpenAIMessage("system", "You are very excited"),
+      OpenAIMessage("user", "")
+      ),
+    Seq(OpenAIMessage("system", "You are very excited")),
     null //scalastyle:ignore null
   ).toDF("messages")
 
@@ -144,11 +149,23 @@ class OpenAIChatCompletionSuite extends TransformerFuzzing[OpenAIChatCompletion]
   }
 
   test("Robustness to bad inputs") {
-    val results = completion.transform(badDf).collect()
-    assert(Option(results.head.getAs[Row](completion.getErrorCol)).isDefined)
-    assert(Option(results.apply(1).getAs[Row](completion.getErrorCol)).isDefined)
-    assert(Option(results.apply(2).getAs[Row](completion.getErrorCol)).isEmpty)
-    assert(Option(results.apply(2).getAs[Row]("out")).isEmpty)
+    val completion: OpenAIChatCompletion = new OpenAIChatCompletion()
+      // by having an invalid deployment name, we should get error for all rows
+      .setDeploymentName("invalid_deployment")
+      .setCustomServiceName(openAIServiceName)
+      .setApiVersion("2023-05-15")
+      .setMaxTokens(5000)
+      .setOutputCol("out")
+      .setMessagesCol("messages")
+      .setTemperature(0)
+      .setSubscriptionKey(openAIAPIKey)
+    val rows = completion.transform(goodDf).collect()
+    assert(rows.length == goodDf.count())
+    rows.foreach { row =>
+      assert(Option(row.getAs[Row]("out")).isEmpty)
+      assert(Option(row.getAs[Row](completion.getErrorCol)).isDefined)
+      assert(Option(row.getAs[Row](completion.getErrorCol)).nonEmpty)
+    }
   }
 
   ignore("Custom EndPoint") {
@@ -174,6 +191,43 @@ class OpenAIChatCompletionSuite extends TransformerFuzzing[OpenAIChatCompletion]
     testCompletion(customEndpointCompletion, goodDf)
   }
 
+  test("shouldSkip method for Good and Bad Data") {
+    shouldSkipRowHelper(goodDf.collect(), expectedToBeSkipped = false)
+    shouldSkipRowHelper(slowDf.collect(), expectedToBeSkipped = false)
+    shouldSkipRowHelper(badDf.collect(), expectedToBeSkipped = true)
+
+    val dfWithSomeEmptyAndNonEmptyUserMessage = Seq(
+      Seq(
+        OpenAIMessage("system", "You are an AI chatbot with red as your favorite color"),
+        OpenAIMessage("user", ""),
+        OpenAIMessage("user", "Whats your favorite color")
+      )
+    ).toDF("messages")
+    shouldSkipRowHelper(dfWithSomeEmptyAndNonEmptyUserMessage.collect(), expectedToBeSkipped = false)
+
+    val dfWithEmptySystemMessage = Seq(
+      Seq(
+        OpenAIMessage("system", ""),
+        OpenAIMessage("user", "Whats your favorite color")
+      ),
+      Seq(
+        OpenAIMessage("system", null), //scalastyle:ignore null
+        OpenAIMessage("user", "Whats your favorite color")
+      ),
+      Seq(OpenAIMessage("user", "Whats your favorite color"))
+    ).toDF("messages")
+    shouldSkipRowHelper(dfWithEmptySystemMessage.collect(), expectedToBeSkipped = false)
+  }
+
+  def shouldSkipRowHelper(rows: Seq[Row], expectedToBeSkipped: Boolean): Unit = {
+    val completion = new OpenAIChatCompletion()
+      .setMessagesCol("messages")
+    for (r <- rows) {
+      assert(completion.shouldSkip(r) == expectedToBeSkipped)
+    }
+  }
+
+
   def testCompletion(completion: OpenAIChatCompletion, df: DataFrame, requiredLength: Int = 10): Unit = {
     val fromRow = ChatCompletionResponse.makeFromRowConverter
     completion.transform(df).collect().foreach(r =>