fix: OpenAIChatCompletion class method shouldSkip method to skip row if no user message

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

@@ -46,8 +46,22 @@ 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" or "assistant" and
+   * "content" as non-empty.
+   */
+  override def shouldSkip(row: Row): Boolean ={
+    val acceptableRoles = Set("user", "assistant")
+    super.shouldSkip(row) || {
+      val messages = Option(row.getAs[Seq[Row]](getMessagesCol)).getOrElse(Seq.empty)
+      messages.isEmpty || !messages
+        // filter out messages system messages
+        .filter(m => acceptableRoles.contains(Option(m.getAs[String]("role")).getOrElse("").toLowerCase))
+        // check if there exists a user message with non-empty content
+        .exists(m => Option(m.getAs[String]("content")).getOrElse("").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,51 @@ 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)
+
+    val dfWithAssistantValidMessages = Seq(
+      Seq(
+        OpenAIMessage("system", "You are an AI chatbot with red as your favorite color"),
+        OpenAIMessage("assistant", "Whats your favorite color")
+      )
+    ).toDF("messages")
+    shouldSkipRowHelper(dfWithAssistantValidMessages.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 =>