split assertVisualAI into assertNoDefectsWithAI and assertWithAI

maestro-ai/src/main/java/maestro/ai/AI.kt

@@ -42,8 +42,8 @@ abstract class AI(
         // * Gemini: https://ai.google.dev/gemini-api/docs/json-mode
 
         val assertVisualSchema: String = run {
-            val resourceStream = this::class.java.getResourceAsStream("/assertVisualAI_schema.json")
-                ?: throw IllegalStateException("Could not find assertVisualAI_schema.json in resources")
+            val resourceStream = this::class.java.getResourceAsStream("/askForDefects_schema.json")
+                ?: throw IllegalStateException("Could not find askForDefects_schema.json in resources")
 
             resourceStream.bufferedReader().use { it.readText() }
         }

maestro-ai/src/main/resources/assertVisualAI_schema.json → maestro-ai/src/main/resources/askForDefects_schema.json

@@ -1,5 +1,5 @@
 {
-  "name": "assertVisualAI",
+  "name": "askForDefects",
   "description": "List of possible defects found in the mobile app's UI",
   "strict": true,
   "schema": {

maestro-orchestra-models/src/main/java/maestro/orchestra/Commands.kt

@@ -364,22 +364,33 @@ data class AssertConditionCommand(
     }
 }
 
-data class AssertVisualAICommand(
-    val assertion: String?,
-    val optional: Boolean = false, /// If true, the command will not fail the flow, but print a warning
+data class AssertNoDefectsWithAICommand(
+    val optional: Boolean = true,
     val label: String? = null,
 ) : Command {
     override fun description(): String {
         if (label != null) return label
 
-        return if (assertion != null) "Assert visual with AI: $assertion"
-            else "Assert visual with AI"
+        return "Assert no defects with AI"
+    }
+
+    override fun evaluateScripts(jsEngine: JsEngine): Command = this
+}
+
+data class AssertWithAICommand(
+    val assertion: String,
+    val optional: Boolean = true,
+    val label: String? = null,
+) : Command {
+    override fun description(): String {
+        if (label != null) return label
+
+        return "Assert no defects with AI: $assertion"
     }
 
     override fun evaluateScripts(jsEngine: JsEngine): Command {
         return copy(
-            assertion = assertion?.evaluateScripts(jsEngine),
-            // TODO: Allow for evaluating script for more properties
+            assertion = assertion.evaluateScripts(jsEngine),
         )
     }
 }

maestro-orchestra-models/src/main/java/maestro/orchestra/MaestroCommand.kt

@@ -36,7 +36,8 @@ data class MaestroCommand(
     val backPressCommand: BackPressCommand? = null,
     @Deprecated("Use assertConditionCommand") val assertCommand: AssertCommand? = null,
     val assertConditionCommand: AssertConditionCommand? = null,
-    val assertVisualAICommand: AssertVisualAICommand? = null,
+    val assertNoDefectsWithAICommand: AssertNoDefectsWithAICommand? = null,
+    val assertWithAICommand: AssertWithAICommand? = null,
     val inputTextCommand: InputTextCommand? = null,
     val inputRandomTextCommand: InputRandomCommand? = null,
     val launchAppCommand: LaunchAppCommand? = null,
@@ -77,7 +78,8 @@ data class MaestroCommand(
         backPressCommand = command as? BackPressCommand,
         assertCommand = command as? AssertCommand,
         assertConditionCommand = command as? AssertConditionCommand,
-        assertVisualAICommand = command as? AssertVisualAICommand,
+        assertNoDefectsWithAICommand = command as? AssertNoDefectsWithAICommand,
+        assertWithAICommand = command as? AssertWithAICommand,
         inputTextCommand = command as? InputTextCommand,
         inputRandomTextCommand = command as? InputRandomCommand,
         launchAppCommand = command as? LaunchAppCommand,
@@ -118,7 +120,8 @@ data class MaestroCommand(
         backPressCommand != null -> backPressCommand
         assertCommand != null -> assertCommand
         assertConditionCommand != null -> assertConditionCommand
-        assertVisualAICommand != null -> assertVisualAICommand
+        assertNoDefectsWithAICommand != null -> assertNoDefectsWithAICommand
+        assertWithAICommand != null -> assertWithAICommand
         inputTextCommand != null -> inputTextCommand
         inputRandomTextCommand != null -> inputRandomTextCommand
         launchAppCommand != null -> launchAppCommand

maestro-orchestra/src/main/java/maestro/orchestra/Orchestra.kt

@@ -23,6 +23,7 @@ import kotlinx.coroutines.runBlocking
 import maestro.*
 import maestro.Filters.asFilter
 import maestro.ai.AI
+import maestro.ai.AI.Companion.AI_KEY_ENV_VAR
 import maestro.ai.Defect
 import maestro.ai.Prediction
 import maestro.ai.antrophic.Claude
@@ -222,7 +223,7 @@ class Orchestra(
     }
 
     private fun initAI(): AI? {
-        val apiKey = System.getenv(AI.AI_KEY_ENV_VAR) ?: return null
+        val apiKey = System.getenv(AI_KEY_ENV_VAR) ?: return null
         val modelName: String? = System.getenv(AI.AI_MODEL_ENV_VAR)
 
         return if (modelName == null) OpenAI(apiKey = apiKey)
@@ -258,7 +259,8 @@ class Orchestra(
             is SwipeCommand -> swipeCommand(command)
             is AssertCommand -> assertCommand(command)
             is AssertConditionCommand -> assertConditionCommand(command)
-            is AssertVisualAICommand -> assertVisualAICommand(command)
+            is AssertNoDefectsWithAICommand -> assertNoDefectsWithAICommand(command)
+            is AssertWithAICommand -> assertWithAICommand(command)
             is InputTextCommand -> inputTextCommand(command)
             is InputRandomCommand -> inputTextRandomCommand(command)
             is LaunchAppCommand -> launchAppCommand(command)
@@ -337,11 +339,43 @@ class Orchestra(
         return false
     }
 
-    private fun assertVisualAICommand(command: AssertVisualAICommand): Boolean = runBlocking {
-        // TODO: make all of Orchestra suspending
+    private fun assertNoDefectsWithAICommand(command: AssertNoDefectsWithAICommand): Boolean = runBlocking {
+        // TODO(bartekpacia): make all of Orchestra suspending
 
         if (ai == null) {
-            throw MaestroException.AINotAvailable("AI is not available")
+            throw MaestroException.AINotAvailable("AI client is not available. Did you export $AI_KEY_ENV_VAR?")
+        }
+
+        val imageData = Buffer()
+        maestro.takeScreenshot(imageData, compressed = false)
+
+        val defects = Prediction.findDefects(
+            aiClient = ai,
+            assertion = null,
+            screen = imageData.copy().readByteArray(),
+            previousFalsePositives = listOf(), // TODO(bartekpacia): take it from WorkspaceConfig (or MaestroConfig?)
+        )
+
+        if (defects.isNotEmpty()) {
+            onCommandGeneratedOutput(command, defects, imageData)
+
+            if (command.optional) throw CommandSkipped
+
+            val word = if (defects.size == 1) "defect" else "defects"
+            throw MaestroException.AssertionFailure(
+                "Ffound ${defects.size} possible $word. See the report after the test completes to learn more.",
+                maestro.viewHierarchy().root,
+            )
+        }
+
+        false
+    }
+
+    private fun assertWithAICommand(command: AssertWithAICommand): Boolean = runBlocking {
+        // TODO(bartekpacia): make all of Orchestra suspending
+
+        if (ai == null) {
+            throw MaestroException.AINotAvailable("AI client is not available. Did you export $AI_KEY_ENV_VAR?")
         }
 
         val imageData = Buffer()

maestro-orchestra/src/main/java/maestro/orchestra/yaml/YamlAssertNoDefectsWithAI.kt

@@ -0,0 +1,6 @@
+package maestro.orchestra.yaml
+
+data class YamlAssertNoDefectsWithAI(
+    val optional: Boolean = true,
+    val label: String? = null,
+)

maestro-orchestra/src/main/java/maestro/orchestra/yaml/YamlAssertVisualAI.kt → maestro-orchestra/src/main/java/maestro/orchestra/yaml/YamlAssertWithAI.kt

@@ -2,18 +2,18 @@ package maestro.orchestra.yaml
 
 import com.fasterxml.jackson.annotation.JsonCreator
 
-data class YamlAssertVisualAI(
-    val assertion: String? = null,
-    val optional: Boolean = false,
+data class YamlAssertWithAI(
+    val assertion: String,
+    val optional: Boolean = true,
     val label: String? = null,
 ) {
 
     companion object {
 
         @JvmStatic
         @JsonCreator(mode = JsonCreator.Mode.DELEGATING)
-        fun parse(assertion: String): YamlAssertVisualAI {
-            return YamlAssertVisualAI(
+        fun parse(assertion: String): YamlAssertWithAI {
+            return YamlAssertWithAI(
                 assertion = assertion,
                 optional = true,
             )

maestro-orchestra/src/main/java/maestro/orchestra/yaml/YamlFluentCommand.kt

@@ -41,7 +41,8 @@ data class YamlFluentCommand(
     val assertVisible: YamlElementSelectorUnion? = null,
     val assertNotVisible: YamlElementSelectorUnion? = null,
     val assertTrue: YamlAssertTrue? = null,
-    val assertVisualAI: YamlAssertVisualAI? = null,
+    val assertNoDefectsWithAI: YamlAssertNoDefectsWithAI? = null,
+    val assertWithAI: YamlAssertWithAI? = null,
     val back: YamlActionBack? = null,
     val clearKeychain: YamlActionClearKeychain? = null,
     val hideKeyboard: YamlActionHideKeyboard? = null,
@@ -116,12 +117,20 @@ data class YamlFluentCommand(
                     )
                 )
             )
-            assertVisualAI != null -> listOf(
+            assertNoDefectsWithAI != null -> listOf(
                 MaestroCommand(
-                    AssertVisualAICommand(
-                        assertion = assertVisualAI.assertion,
-                        optional = assertVisualAI.optional,
-                        label = assertVisualAI.label,
+                    AssertNoDefectsWithAICommand(
+                        optional = assertNoDefectsWithAI.optional,
+                        label = assertNoDefectsWithAI.label,
+                    )
+                )
+            )
+            assertWithAI != null -> listOf(
+                MaestroCommand(
+                    AssertWithAICommand(
+                        assertion = assertWithAI.assertion,
+                        optional = assertWithAI.optional,
+                        label = assertWithAI.label,
                     )
                 )
             )
@@ -698,8 +707,8 @@ data class YamlFluentCommand(
                     toggleAirplaneMode = YamlToggleAirplaneMode()
                 )
 
-                "assertVisualAI" -> YamlFluentCommand(
-                    assertVisualAI = YamlAssertVisualAI()
+                "assertNoDefectsWithAI" -> YamlFluentCommand(
+                    assertNoDefectsWithAI = YamlAssertNoDefectsWithAI()
                 )
 
                 else -> throw SyntaxError("Invalid command: \"$stringCommand\"")