Skip to content

Commit

Permalink
saving LLM output: text done, screenshots WIP
Browse files Browse the repository at this point in the history
  • Loading branch information
bartekpacia committed Aug 14, 2024
1 parent 644c94c commit 65fd303
Show file tree
Hide file tree
Showing 5 changed files with 117 additions and 60 deletions.
29 changes: 19 additions & 10 deletions maestro-cli/src/main/java/maestro/cli/report/TestDebugReporter.kt
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import com.fasterxml.jackson.annotation.JsonInclude
import com.fasterxml.jackson.annotation.JsonProperty
import com.fasterxml.jackson.databind.JsonMappingException
import com.fasterxml.jackson.databind.ObjectMapper
import maestro.Driver
import maestro.MaestroException
import maestro.TreeNode
import maestro.cli.runner.CommandStatus
Expand All @@ -27,7 +26,6 @@ import java.time.LocalDateTime
import java.time.format.DateTimeFormatter
import java.time.temporal.ChronoUnit
import java.util.IdentityHashMap
import java.util.Properties
import kotlin.io.path.absolutePathString
import kotlin.io.path.exists

Expand All @@ -48,7 +46,9 @@ object TestDebugReporter {
mapper.setSerializationInclusion(JsonInclude.Include.NON_EMPTY)
}

fun saveFlow(flowName: String, debugOutput: FlowDebugOutput, aiOutput: AIOutput, path: Path) {
fun saveFlow(flowName: String, debugOutput: FlowDebugOutput, aiOutput: FlowAIOutput?, path: Path) {
// TODO(bartekpacia): Potentially accept a single "FlowPersistentOutput object

println("TestDebugReporter.saveFlow: saving flow metadata to $path")

// commands
Expand All @@ -72,11 +72,18 @@ object TestDebugReporter {
CommandStatus.FAILED -> ""
else -> ""
}
val name = "screenshot-$status-${it.timestamp}-(${flowName}).png"
val file = File(path.absolutePathString(), name)
val filename = "screenshot-$status-${it.timestamp}-(${flowName}).png"
val file = File(path.absolutePathString(), filename)

it.screenshot.copyTo(file)
}

// AI output
aiOutput?.let {
val filename = "ai-(${flowName.replace("/", "_")}).json"
val file = File(path.absolutePathString(), filename)
mapper.writeValue(file, it)
}
}

fun deleteOldFiles(days: Long = 14) {
Expand Down Expand Up @@ -139,7 +146,7 @@ object TestDebugReporter {
return debugOutput
}

fun buildDefaultDebugOutputPath(debugRootPath: String): Path {
private fun buildDefaultDebugOutputPath(debugRootPath: String): Path {
val preamble = arrayOf(".maestro", "tests")
val foldername = DateTimeFormatter.ofPattern("yyyy-MM-dd_HHmmss").format(LocalDateTime.now())
return Paths.get(debugRootPath, *preamble, foldername)
Expand Down Expand Up @@ -171,17 +178,19 @@ data class FlowDebugOutput(
var exception: MaestroException? = null,
) {
data class Screenshot(
val screenshot: File, val timestamp: Long, val status: CommandStatus
val screenshot: File,
val timestamp: Long,
val status: CommandStatus,
)
}

data class FlowAIOutput(
@JsonProperty("flow_name") var flowName: String,
@JsonProperty("flow_file_path") var flowFilePath: String,
@JsonProperty("flow_name") val flowName: String,
@JsonProperty("flow_file_path") val flowFilePath: String,
val outputs: MutableList<AIOutput> = mutableListOf(),
)

data class AIOutput(
@JsonProperty("screenshot_path") val screenshotPath: File,
val defect: List<Defect>,
val defects: List<Defect>,
)
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ package maestro.cli.runner
import maestro.Maestro
import maestro.MaestroException
import maestro.cli.device.Device
import maestro.cli.report.AIOutput
import maestro.cli.report.CommandDebugMetadata
import maestro.cli.report.FlowAIOutput
import maestro.cli.report.FlowDebugOutput
import maestro.cli.runner.resultview.ResultView
import maestro.cli.runner.resultview.UiState
Expand All @@ -33,6 +35,7 @@ import maestro.orchestra.Orchestra
import maestro.orchestra.OrchestraAppState
import maestro.orchestra.yaml.YamlCommandReader
import maestro.utils.Insight
import okio.Buffer
import org.slf4j.LoggerFactory
import java.io.File
import java.util.IdentityHashMap
Expand All @@ -46,7 +49,8 @@ object MaestroCommandRunner {
device: Device?,
view: ResultView,
commands: List<MaestroCommand>,
debug: FlowDebugOutput
debugOutput: FlowDebugOutput,
aiOutput: FlowAIOutput,
): Result {
val config = YamlCommandReader.getConfig(commands)
val initFlow = config?.initFlow
Expand All @@ -56,23 +60,20 @@ object MaestroCommandRunner {
val commandStatuses = IdentityHashMap<MaestroCommand, CommandStatus>()
val commandMetadata = IdentityHashMap<MaestroCommand, Orchestra.CommandMetadata>()

// debug
val debugCommands = debug.commands
val debugScreenshots = debug.screenshots

fun takeDebugScreenshot(status: CommandStatus): File? {
val containsFailed = debugScreenshots.any { it.status == CommandStatus.FAILED }
val containsFailed = debugOutput.screenshots.any { it.status == CommandStatus.FAILED }

// Avoids duplicate failed images from parent commands
if (containsFailed && status == CommandStatus.FAILED) {
return null
}

val result = kotlin.runCatching {
val out = File.createTempFile("screenshot-${System.currentTimeMillis()}", ".png")
val out = File
.createTempFile("screenshot-${System.currentTimeMillis()}", ".png")
.also { it.deleteOnExit() } // save to another dir before exiting
maestro.takeScreenshot(out, false)
debugScreenshots.add(
debugOutput.screenshots.add(
FlowDebugOutput.Screenshot(
screenshot = out,
timestamp = System.currentTimeMillis(),
Expand All @@ -85,6 +86,12 @@ object MaestroCommandRunner {
return result.getOrNull()
}

fun writeAIscreenshot(buffer: Buffer): File {
val out = File.createTempFile("ai-screenshot-${System.currentTimeMillis()}", ".png")
out.outputStream().use { it.write(buffer.readByteArray()) }
return out
}

fun refreshUi() {
view.setState(
UiState.Running(
Expand Down Expand Up @@ -120,7 +127,7 @@ object MaestroCommandRunner {
onCommandStart = { _, command ->
logger.info("${command.description()} RUNNING")
commandStatuses[command] = CommandStatus.RUNNING
debugCommands[command] = CommandDebugMetadata(
debugOutput.commands[command] = CommandDebugMetadata(
timestamp = System.currentTimeMillis(),
status = CommandStatus.RUNNING
)
Expand All @@ -130,25 +137,25 @@ object MaestroCommandRunner {
onCommandComplete = { _, command ->
logger.info("${command.description()} COMPLETED")
commandStatuses[command] = CommandStatus.COMPLETED
debugCommands[command]?.let {
it.status = CommandStatus.COMPLETED
it.calculateDuration()
debugOutput.commands[command]?.apply {
status = CommandStatus.COMPLETED
calculateDuration()
}
refreshUi()
},
onCommandFailed = { _, command, e ->
debugCommands[command]?.let {
it.status = CommandStatus.FAILED
it.calculateDuration()
it.error = e
debugOutput.commands[command]?.apply {
status = CommandStatus.FAILED
calculateDuration()
error = e
}

takeDebugScreenshot(CommandStatus.FAILED)

if (e !is MaestroException) {
throw e
} else {
debug.exception = e
debugOutput.exception = e
}

logger.info("${command.description()} FAILED")
Expand All @@ -159,16 +166,16 @@ object MaestroCommandRunner {
onCommandSkipped = { _, command ->
logger.info("${command.description()} SKIPPED")
commandStatuses[command] = CommandStatus.SKIPPED
debugCommands[command]?.let {
it.status = CommandStatus.SKIPPED
debugOutput.commands[command]?.apply {
status = CommandStatus.SKIPPED
}
refreshUi()
},
onCommandReset = { command ->
logger.info("${command.description()} PENDING")
commandStatuses[command] = CommandStatus.PENDING
debugCommands[command]?.let {
it.status = CommandStatus.PENDING
debugOutput.commands[command]?.apply {
status = CommandStatus.PENDING
}
refreshUi()
},
Expand All @@ -177,6 +184,16 @@ object MaestroCommandRunner {
commandMetadata[command] = metadata
refreshUi()
},
onCommandGeneratedOutput = { command, defects, screenshot ->
logger.info("${command.description()} OUTPUTTED")
val screenshotPath = writeAIscreenshot(screenshot)
aiOutput.outputs.add(
AIOutput(
screenshotPath = screenshotPath,
defects = defects
)
)
}
)

val flowSuccess = orchestra.runFlow(commands)
Expand Down
17 changes: 15 additions & 2 deletions maestro-cli/src/main/java/maestro/cli/runner/TestRunner.kt
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,19 @@ object TestRunner {
debugOutputPath: Path
): Int {
val debugOutput = FlowDebugOutput()
val aiOutput = FlowAIOutput()
var aiOutput = FlowAIOutput(
flowName = flowFile.nameWithoutExtension,
flowFilePath = flowFile.absolutePath,
)

val result = runCatching(resultView, maestro) {
val commands = YamlCommandReader.readCommands(flowFile.toPath())
.withEnv(env)

YamlCommandReader.getConfig(commands)?.name?.let {
aiOutput = aiOutput.copy(flowName = it)
}

MaestroCommandRunner.runCommands(
maestro,
device,
Expand Down Expand Up @@ -112,7 +120,12 @@ object TestRunner {
device,
resultView,
commands,
FlowDebugOutput()
FlowDebugOutput(),
// TODO: bartekpacia - make AI outputs work in continuous mode
FlowAIOutput(
flowName = "TODO",
flowFilePath = flowFile.absolutePath,
),
)
}.get()
}
Expand Down
Loading

0 comments on commit 65fd303

Please sign in to comment.