From 8fc42e96b24e0009a26fab2e996da3c60ae72ce1 Mon Sep 17 00:00:00 2001 From: Jingxiao Zheng Date: Wed, 22 Jul 2026 14:36:34 -0700 Subject: [PATCH] Implement ConversationManager for session history tracking. PiperOrigin-RevId: 952337032 --- Android/src/app/eval/AndroidManifest.xml | 47 +++ .../edge/gallery/eval/ConversationManager.kt | 70 +++++ .../ai/edge/gallery/eval/EvalReceiver.kt | 49 +++ .../ai/edge/gallery/eval/EvalService.kt | 91 ++++++ .../ai/edge/gallery/eval/MiniHttpServer.kt | 153 ++++++++++ .../ai/edge/gallery/eval/ModelManager.kt | 195 ++++++++++++ .../ai/edge/gallery/eval/PromptParser.kt | 145 +++++++++ .../gallery/eval/ConversationManagerTest.kt | 159 ++++++++++ .../ai/edge/gallery/eval/EvalAppTestSuite.kt | 31 ++ .../ai/edge/gallery/eval/EvalReceiverTest.kt | 53 ++++ .../ai/edge/gallery/eval/EvalServiceTest.kt | 137 +++++++++ .../edge/gallery/eval/MiniHttpServerTest.kt | 105 +++++++ .../ai/edge/gallery/eval/ModelManagerTest.kt | 284 ++++++++++++++++++ .../ai/edge/gallery/eval/PromptParserTest.kt | 158 ++++++++++ 14 files changed, 1677 insertions(+) create mode 100644 Android/src/app/eval/AndroidManifest.xml create mode 100644 Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ConversationManager.kt create mode 100644 Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalReceiver.kt create mode 100644 Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalService.kt create mode 100644 Android/src/app/eval/java/com/google/ai/edge/gallery/eval/MiniHttpServer.kt create mode 100644 Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ModelManager.kt create mode 100644 Android/src/app/eval/java/com/google/ai/edge/gallery/eval/PromptParser.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ConversationManagerTest.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalAppTestSuite.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalReceiverTest.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalServiceTest.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/MiniHttpServerTest.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ModelManagerTest.kt create mode 100644 Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/PromptParserTest.kt diff --git a/Android/src/app/eval/AndroidManifest.xml b/Android/src/app/eval/AndroidManifest.xml new file mode 100644 index 000000000..bdb2a00c1 --- /dev/null +++ b/Android/src/app/eval/AndroidManifest.xml @@ -0,0 +1,47 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ConversationManager.kt b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ConversationManager.kt new file mode 100644 index 000000000..a8078b856 --- /dev/null +++ b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ConversationManager.kt @@ -0,0 +1,70 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.util.Log +import com.google.ai.edge.gallery.data.Model +import com.google.ai.edge.gallery.runtime.LlmModelHelper +import java.util.concurrent.ConcurrentHashMap + +class ConversationManager { + + private val modelHistory = ConcurrentHashMap>() + + fun checkHistoryAndReset( + modelName: String, + incomingHistory: List, + model: Model, + helper: LlmModelHelper, + ) { + val cachedHistory = modelHistory[modelName] + if (cachedHistory == null || cachedHistory != incomingHistory) { + Log.i(TAG, "History mismatch or new session. Resetting conversation for $modelName.") + val lmHistory = PromptParser.convertToLmMessages(incomingHistory) + helper.resetConversation( + model = model, + supportImage = model.llmSupportImage, + supportAudio = model.llmSupportAudio, + initialMessages = lmHistory, + ) + modelHistory[modelName] = incomingHistory + } else { + Log.i(TAG, "History match. Appending to existing conversation for $modelName.") + } + } + + fun appendTurn(modelName: String, promptContentStr: String, assistantResult: String) { + val currentHistory = modelHistory[modelName] ?: emptyList() + modelHistory[modelName] = + currentHistory + + HistoryMessage("user", promptContentStr) + + HistoryMessage("assistant", assistantResult) + } + + // Exposed for testing + fun getHistory(modelName: String): List? { + return modelHistory[modelName] + } + + // Exposed for testing + fun updateHistory(modelName: String, history: List) { + modelHistory[modelName] = history + } + + companion object { + private const val TAG = "ConversationManager" + } +} diff --git a/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalReceiver.kt b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalReceiver.kt new file mode 100644 index 000000000..44e9121c0 --- /dev/null +++ b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalReceiver.kt @@ -0,0 +1,49 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.content.BroadcastReceiver +import android.content.Context +import android.content.Intent +import android.util.Log + +class EvalReceiver : BroadcastReceiver() { + override fun onReceive(context: Context, intent: Intent) { + val action = intent.action + Log.i(TAG, "Received action: $action") + val serviceIntent = + Intent(context, EvalService::class.java).apply { + this.action = action + if (intent.hasExtra("port")) { + putExtra("port", intent.getIntExtra("port", 8080)) + } + } + + if (action == EvalService.ACTION_STOP_SERVER) { + context.stopService(Intent(context, EvalService::class.java)) + } else { + try { + context.startForegroundService(serviceIntent) + } catch (e: Exception) { + Log.e(TAG, "Failed to start service", e) + } + } + } + + companion object { + private const val TAG = "EvalReceiver" + } +} diff --git a/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalService.kt b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalService.kt new file mode 100644 index 000000000..4683c7dc3 --- /dev/null +++ b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/EvalService.kt @@ -0,0 +1,91 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.app.Notification +import android.app.NotificationChannel +import android.app.NotificationManager +import android.app.Service +import android.content.Intent +import android.content.pm.ServiceInfo +import android.os.Build +import android.os.IBinder +import android.util.Log +import androidx.core.app.NotificationCompat + +class EvalService : Service() { + + override fun onCreate() { + super.onCreate() + createNotificationChannel() + } + + override fun onStartCommand(intent: Intent?, flags: Int, startId: Int): Int { + val action = intent?.action + Log.i(TAG, "onStartCommand action: $action") + + if (action == ACTION_STOP_SERVER) { + stopSelf() + return START_NOT_STICKY + } + + // Extract optional model configuration from the intent. + // In CL 3, these will be used to pre-initialize the model. + val port = intent?.getIntExtra("port", 8080) ?: 8080 + + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.UPSIDE_DOWN_CAKE) { + startForeground( + NOTIFICATION_ID, + createNotification(port), + ServiceInfo.FOREGROUND_SERVICE_TYPE_SPECIAL_USE, + ) + } else { + startForeground(NOTIFICATION_ID, createNotification(port)) + } + + return START_STICKY + } + + override fun onDestroy() { + super.onDestroy() + } + + override fun onBind(intent: Intent?): IBinder? = null + + private fun createNotificationChannel() { + val channel = + NotificationChannel(CHANNEL_ID, "Eval Service Channel", NotificationManager.IMPORTANCE_LOW) + val manager = getSystemService(NotificationManager::class.java) + manager?.createNotificationChannel(channel) + } + + private fun createNotification(port: Int): Notification { + return NotificationCompat.Builder(this, CHANNEL_ID) + .setContentTitle("Gallery Eval Server") + .setContentText("Running on port $port") + .setSmallIcon(android.R.drawable.ic_media_play) + .build() + } + + companion object { + private const val TAG = "EvalService" + private const val CHANNEL_ID = "EvalServiceChannel" + private const val NOTIFICATION_ID = 1 + + const val ACTION_START_SERVER = "com.google.ai.edge.gallery.eval.START_SERVER" + const val ACTION_STOP_SERVER = "com.google.ai.edge.gallery.eval.STOP_SERVER" + } +} diff --git a/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/MiniHttpServer.kt b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/MiniHttpServer.kt new file mode 100644 index 000000000..5dbcfa5b7 --- /dev/null +++ b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/MiniHttpServer.kt @@ -0,0 +1,153 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.util.Log +import java.io.BufferedInputStream +import java.io.OutputStream +import java.net.ServerSocket +import java.net.Socket +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.cancel +import kotlinx.coroutines.launch + +class MiniHttpServer(val port: Int, val handler: (Request) -> Response) { + private var serverSocket: ServerSocket? = null + private val scope = CoroutineScope(Dispatchers.IO) + @Volatile private var running = false + + fun start() { + running = true + serverSocket = ServerSocket(port) + scope.launch { + while (running) { + try { + val socket = serverSocket?.accept() ?: break + scope.launch { handleConnection(socket) } + } catch (e: Exception) { + if (running) Log.e(TAG, "Error accepting connection", e) + } + } + } + } + + fun stop() { + running = false + serverSocket?.close() + serverSocket = null + scope.cancel() + } + + private fun handleConnection(socket: Socket) { + try { + val input = BufferedInputStream(socket.inputStream) + val output = socket.outputStream + + val headerBuilder = StringBuilder() + var currentByte = input.read() + + // Read headers until \r\n\r\n + while (currentByte != -1) { + headerBuilder.append(currentByte.toChar()) + if (headerBuilder.endsWith("\r\n\r\n")) break + currentByte = input.read() + } + + val headerString = headerBuilder.toString() + val headerLines = headerString.lines() + if (headerLines.isEmpty()) return + + val requestLineParts = headerLines[0].split(" ") + if (requestLineParts.size < 3) return + val method = requestLineParts[0] + val path = requestLineParts[1] + + // Parse headers + val headers = mutableMapOf() + var contentLength = 0 + for (i in 1 until headerLines.size) { + val line = headerLines[i] + if (line.isEmpty()) continue + val parts = line.split(":", limit = 2) + if (parts.size == 2) { + val key = parts[0].trim().lowercase() + val value = parts[1].trim() + headers[key] = value + if (key == "content-length") { + contentLength = value.toIntOrNull() ?: 0 + } + } + } + + // Safely read exact number of BYTES for the body + val body = + if (contentLength > 0) { + val bodyBytes = ByteArray(contentLength) + var bytesRead = 0 + while (bytesRead < contentLength) { + val read = input.read(bodyBytes, bytesRead, contentLength - bytesRead) + if (read == -1) break + bytesRead += read + } + String(bodyBytes, Charsets.UTF_8) // Decode to string safely here! + } else "" + + val request = Request(method, path, headers, body) + val response = handler(request) + writeResponse(output, response) + } catch (e: Exception) { + Log.e(TAG, "Error handling connection", e) + } finally { + socket.close() + } + } + + private fun writeResponse(output: OutputStream, response: Response) { + // Use the byte size of the encoded UTF-8 body to compute Content-Length. + // Using string character length (response.body.length) would cause response + // truncation on the client side for responses containing multi-byte characters. + val bodyBytes = response.body.toByteArray() + val statusLine = "HTTP/1.1 ${response.status.code} ${response.status.message}\r\n" + output.write(statusLine.toByteArray()) + output.write("Content-Type: ${response.contentType}\r\n".toByteArray()) + output.write("Content-Length: ${bodyBytes.size}\r\n".toByteArray()) + output.write("Connection: close\r\n".toByteArray()) + output.write("\r\n".toByteArray()) + output.write(bodyBytes) + output.flush() + } + + data class Request( + val method: String, + val path: String, + val headers: Map, + val body: String, + ) + + data class Response(val status: Status, val contentType: String, val body: String) + + enum class Status(val code: Int, val message: String) { + OK(200, "OK"), + BAD_REQUEST(400, "Bad Request"), + NOT_FOUND(404, "Not Found"), + INTERNAL_ERROR(500, "Internal Server Error"), + } + + companion object { + private const val TAG = "MiniHttpServer" + } +} diff --git a/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ModelManager.kt b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ModelManager.kt new file mode 100644 index 000000000..424c2d1f9 --- /dev/null +++ b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/ModelManager.kt @@ -0,0 +1,195 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.content.Context +import android.util.Log +import com.google.ai.edge.gallery.data.ConfigKeys +import com.google.ai.edge.gallery.data.Model +import com.google.ai.edge.gallery.data.RuntimeType +import com.google.ai.edge.gallery.runtime.LlmModelHelper +import com.google.ai.edge.gallery.runtime.runtimeHelper +import java.util.concurrent.ConcurrentHashMap +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch + +open class ModelManager( + private val context: Context, + private val coroutineScope: CoroutineScope = CoroutineScope(Dispatchers.Default), +) { + + private val modelCache = ConcurrentHashMap() + private val initJobs = ConcurrentHashMap>() + + /** + * Pre-initializes the model in the background. This ensures the model is cached with the correct + * capabilities (e.g. vision) before the first HTTP request arrives. + */ + open fun preInitModel( + modelPath: String, + supportImage: Boolean, + supportAudio: Boolean, + accelerator: String, + ) { + Log.i( + TAG, + "Pre-initializing model $modelPath (image: $supportImage, audio: $supportAudio, acc: $accelerator)", + ) + val helper = getHelperForModelName(modelPath) + if (helper == null) { + Log.e(TAG, "Failed to get helper for model $modelPath") + return + } + coroutineScope.launch { + val unused = getOrInitModel(modelPath, helper, supportImage, supportAudio, accelerator) + } + } + + suspend fun getOrInitModel( + modelName: String, + helper: LlmModelHelper, + supportImage: Boolean, + supportAudio: Boolean, + accelerator: String, + ): Model? { + val job: CompletableDeferred + var isNew = false + + // Synchronize on modelCache to make the check-and-register atomic. + // This prevents a race condition where a background pre-initialization + // finishes and removes itself from initJobs before a concurrent request + // thread can find it in either modelCache or initJobs. + synchronized(modelCache) { + val cachedModel = modelCache[modelName] + if (cachedModel != null) { + return cachedModel + } + val existingJob = initJobs[modelName] + if (existingJob != null) { + job = existingJob + } else { + job = CompletableDeferred() + initJobs[modelName] = job + isNew = true + } + } + + if (isNew) { + val model = + resolveModel(modelName, supportImage, supportAudio, accelerator) + ?: run { + synchronized(modelCache) { initJobs.remove(modelName)?.let {} } + job.completeExceptionally(IllegalArgumentException("Unknown model: $modelName")) + return null + } + + Log.i(TAG, "Initializing model: $modelName") + + helper.initialize( + context = context, + model = model, + taskId = "eval_task", + supportImage = model.llmSupportImage, + supportAudio = model.llmSupportAudio, + coroutineScope = coroutineScope, + onDone = { status -> + Log.i(TAG, "Init status for $modelName: $status") + val isSuccess = + status == "" || status == "Feature is available" || status == "Download completed" + val isDownloading = status.startsWith("Downloading") + + if (isSuccess) { + synchronized(modelCache) { + modelCache[modelName] = model + initJobs.remove(modelName)?.let {} + } + job.complete(model) + } else if (!isDownloading) { + synchronized(modelCache) { initJobs.remove(modelName)?.let {} } + job.completeExceptionally(RuntimeException("Init failed: $status")) + } + }, + ) + } + + return try { + job.await() + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + Log.e(TAG, "Failed to initialize model $modelName", e) + null + } + } + + open fun resolveModel( + modelName: String, + supportImage: Boolean = false, + supportAudio: Boolean = false, + accelerator: String = "CPU", + ): Model? { + val configValues = + mapOf( + ConfigKeys.ACCELERATOR.label to accelerator, + ConfigKeys.VISION_ACCELERATOR.label to accelerator, + ) + return when { + modelName.contains("aicore", ignoreCase = true) -> { + Model( + name = modelName, + runtimeType = RuntimeType.AICORE, + isLlm = true, + llmSupportImage = supportImage, + llmSupportAudio = supportAudio, + ) + .apply { this.configValues = configValues } + } + modelName.startsWith("/") -> { + Model( + name = modelName, + runtimeType = RuntimeType.LITERT_LM, + localModelFilePathOverride = modelName, + isLlm = true, + llmSupportImage = supportImage, + llmSupportAudio = supportAudio, + ) + .apply { this.configValues = configValues } + } + else -> { + Model( + name = modelName, + runtimeType = RuntimeType.LITERT_LM, + isLlm = true, + llmSupportImage = supportImage, + llmSupportAudio = supportAudio, + ) + .apply { this.configValues = configValues } + } + } + } + + open fun getHelperForModelName(modelName: String): LlmModelHelper? { + val model = resolveModel(modelName) ?: return null + return model.runtimeHelper + } + + companion object { + private const val TAG = "ModelManager" + } +} diff --git a/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/PromptParser.kt b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/PromptParser.kt new file mode 100644 index 000000000..efdf174c9 --- /dev/null +++ b/Android/src/app/eval/java/com/google/ai/edge/gallery/eval/PromptParser.kt @@ -0,0 +1,145 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.graphics.Bitmap +import android.graphics.BitmapFactory +import android.util.Base64 +import android.util.Log +import com.google.ai.edge.litertlm.Content as LmContent +import com.google.ai.edge.litertlm.Contents as LmContents +import com.google.ai.edge.litertlm.Message as LmMessage +import org.json.JSONArray + +object PromptParser { + private const val TAG = "PromptParser" + + data class ParsedPrompt( + val text: String, + val images: List, + val audioClips: List, + ) + + fun parseContent(contentVal: Any): ParsedPrompt { + val images = mutableListOf() + val audioClips = mutableListOf() + var textPrompt = "" + + if (contentVal is JSONArray) { + for (i in 0 until contentVal.length()) { + val part = contentVal.getJSONObject(i) + val type = part.getString("type") + when (type) { + "text" -> { + textPrompt += part.getString("text") + } + "image_url" -> { + val urlObj = part.getJSONObject("image_url") + val url = urlObj.getString("url") + if (url.startsWith("data:image/")) { + val base64Data = url.substringAfter("base64,") + try { + val decodedString = Base64.decode(base64Data, Base64.DEFAULT) + val bitmap = BitmapFactory.decodeByteArray(decodedString, 0, decodedString.size) + if (bitmap != null) { + images.add(bitmap) + } else { + Log.e(TAG, "Failed to decode bitmap from base64") + } + } catch (e: Exception) { + Log.e(TAG, "Failed to parse base64 image", e) + } + } + } + "input_audio" -> { + val audioObj = part.getJSONObject("input_audio") + val base64Data = audioObj.getString("data") + try { + val decodedString = Base64.decode(base64Data, Base64.DEFAULT) + audioClips.add(decodedString) + } catch (e: Exception) { + Log.e(TAG, "Failed to parse base64 audio", e) + } + } + } + } + } else { + textPrompt = contentVal.toString() + } + return ParsedPrompt(textPrompt, images, audioClips) + } + + fun jsonToContents(contentVal: Any): LmContents { + if (contentVal is JSONArray) { + val parts = mutableListOf() + for (i in 0 until contentVal.length()) { + val partObj = contentVal.getJSONObject(i) + val type = partObj.getString("type") + when (type) { + "text" -> { + parts.add(LmContent.Text(partObj.getString("text"))) + } + "image_url" -> { + val urlObj = partObj.getJSONObject("image_url") + val url = urlObj.getString("url") + if (url.startsWith("data:image/")) { + val base64Data = url.substringAfter("base64,") + try { + val decodedString = Base64.decode(base64Data, Base64.DEFAULT) + parts.add(LmContent.ImageBytes(decodedString)) + } catch (e: Exception) { + Log.e(TAG, "Failed to decode base64 image in jsonToContents", e) + } + } + } + "input_audio" -> { + val audioObj = partObj.getJSONObject("input_audio") + val base64Data = audioObj.getString("data") + try { + val decodedString = Base64.decode(base64Data, Base64.DEFAULT) + parts.add(LmContent.AudioBytes(decodedString)) + } catch (e: Exception) { + Log.e(TAG, "Failed to decode base64 audio in jsonToContents", e) + } + } + } + } + return LmContents.of(parts) + } else { + return LmContents.of(contentVal.toString()) + } + } + + fun convertToLmMessages(history: List): List { + return history.mapNotNull { msg -> + val contentVal = + try { + if (msg.content.startsWith("[")) JSONArray(msg.content) else msg.content + } catch (e: Exception) { + msg.content + } + val lmContents = jsonToContents(contentVal) + when (msg.role) { + "user" -> LmMessage.user(lmContents) + "assistant" -> LmMessage.model(lmContents) + "system" -> LmMessage.system(lmContents) + else -> null + } + } + } +} + +data class HistoryMessage(val role: String, val content: String) diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ConversationManagerTest.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ConversationManagerTest.kt new file mode 100644 index 000000000..2666310af --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ConversationManagerTest.kt @@ -0,0 +1,159 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.content.Context +import android.graphics.Bitmap +import androidx.test.core.app.ApplicationProvider +import com.google.ai.edge.gallery.data.Model +import com.google.ai.edge.gallery.data.RuntimeType +import com.google.ai.edge.gallery.runtime.CleanUpListener +import com.google.ai.edge.gallery.runtime.LlmModelHelper +import com.google.ai.edge.gallery.runtime.ResultListener +import com.google.ai.edge.litertlm.Contents +import com.google.ai.edge.litertlm.Message +import com.google.ai.edge.litertlm.ToolProvider +import com.google.common.truth.Truth.assertThat +import java.util.concurrent.atomic.AtomicInteger +import kotlinx.coroutines.CoroutineScope +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner + +@RunWith(RobolectricTestRunner::class) +class ConversationManagerTest { + + private lateinit var context: Context + private lateinit var conversationManager: ConversationManager + private lateinit var model: Model + private lateinit var helper: FakeLlmModelHelper + private val modelName = "test-model" + + @Before + fun setUp() { + context = ApplicationProvider.getApplicationContext() + conversationManager = ConversationManager() + model = Model(name = modelName, runtimeType = RuntimeType.LITERT_LM, isLlm = true) + helper = FakeLlmModelHelper() + } + + @Test + fun checkHistoryAndReset_newSession_resetsAndCaches() { + val history = listOf(HistoryMessage("user", "Hello")) + + conversationManager.checkHistoryAndReset(modelName, history, model, helper) + + assertThat(helper.resetCallCount.get()).isEqualTo(1) + assertThat(helper.lastInitialMessagesSize).isEqualTo(1) + assertThat(conversationManager.getHistory(modelName)).isEqualTo(history) + } + + @Test + fun checkHistoryAndReset_matchingHistory_doesNotReset() { + val history = listOf(HistoryMessage("user", "Hello")) + + // First call to cache it + conversationManager.checkHistoryAndReset(modelName, history, model, helper) + assertThat(helper.resetCallCount.get()).isEqualTo(1) + + // Second call with same history + conversationManager.checkHistoryAndReset(modelName, history, model, helper) + assertThat(helper.resetCallCount.get()).isEqualTo(1) // Should still be 1 + } + + @Test + fun checkHistoryAndReset_mismatchingHistory_resetsAndUpdates() { + val history1 = listOf(HistoryMessage("user", "Hello")) + val history2 = + listOf( + HistoryMessage("user", "Hello"), + HistoryMessage("assistant", "Hi"), + HistoryMessage("user", "How are you?"), + ) + + // First call + conversationManager.checkHistoryAndReset(modelName, history1, model, helper) + assertThat(helper.resetCallCount.get()).isEqualTo(1) + + // Second call with different history + conversationManager.checkHistoryAndReset(modelName, history2, model, helper) + assertThat(helper.resetCallCount.get()).isEqualTo(2) + assertThat(helper.lastInitialMessagesSize).isEqualTo(3) + assertThat(conversationManager.getHistory(modelName)).isEqualTo(history2) + } + + @Test + fun appendTurn_appendsCorrectly() { + val initialHistory = listOf(HistoryMessage("user", "Hello")) + conversationManager.updateHistory(modelName, initialHistory) + + conversationManager.appendTurn(modelName, "How are you?", "I am fine") + + val expected = + initialHistory + + HistoryMessage("user", "How are you?") + + HistoryMessage("assistant", "I am fine") + assertThat(conversationManager.getHistory(modelName)).isEqualTo(expected) + } + + class FakeLlmModelHelper : LlmModelHelper { + val resetCallCount = AtomicInteger(0) + var lastInitialMessagesSize = -1 + + override fun initialize( + context: Context, + model: Model, + taskId: String, + supportImage: Boolean, + supportAudio: Boolean, + onDone: (String) -> Unit, + systemInstruction: Contents?, + tools: List, + enableConversationConstrainedDecoding: Boolean, + coroutineScope: CoroutineScope?, + ) {} + + override fun resetConversation( + model: Model, + supportImage: Boolean, + supportAudio: Boolean, + systemInstruction: Contents?, + tools: List, + enableConversationConstrainedDecoding: Boolean, + initialMessages: List, + ) { + resetCallCount.incrementAndGet() + lastInitialMessagesSize = initialMessages.size + } + + override fun cleanUp(model: Model, onDone: () -> Unit) {} + + override fun runInference( + model: Model, + input: String, + resultListener: ResultListener, + cleanUpListener: CleanUpListener, + onError: (message: String) -> Unit, + images: List, + audioClips: List, + coroutineScope: CoroutineScope?, + extraContext: Map?, + ) {} + + override fun stopResponse(model: Model) {} + } +} diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalAppTestSuite.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalAppTestSuite.kt new file mode 100644 index 000000000..f25f91f79 --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalAppTestSuite.kt @@ -0,0 +1,31 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import org.junit.runner.RunWith +import org.junit.runners.Suite + +/** Test suite to run all local unit tests for the on-device evaluation app. */ +@RunWith(Suite::class) +@Suite.SuiteClasses( + MiniHttpServerTest::class, + EvalServiceTest::class, + EvalReceiverTest::class, + PromptParserTest::class, + ModelManagerTest::class, + ConversationManagerTest::class, +) +class EvalAppTestSuite diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalReceiverTest.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalReceiverTest.kt new file mode 100644 index 000000000..ab193d1a9 --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalReceiverTest.kt @@ -0,0 +1,53 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.app.Application +import android.content.Intent +import androidx.test.core.app.ApplicationProvider +import com.google.common.truth.Truth.assertWithMessage +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner +import org.robolectric.Shadows + +/** Unit tests for [EvalReceiver] to verify broadcast routing to [EvalService]. */ +@RunWith(RobolectricTestRunner::class) +class EvalReceiverTest { + + @Test + fun onReceive_forwardsIntentToService() { + val context = ApplicationProvider.getApplicationContext() + val receiver = EvalReceiver() + + val intent = Intent(EvalService.ACTION_START_SERVER).apply { putExtra("port", 9090) } + + receiver.onReceive(context, intent) + + val shadowApp = Shadows.shadowOf(context) + val startedIntent: Intent? = shadowApp.nextStartedService + + assertWithMessage("startedIntent").that(startedIntent).isNotNull() + val nonNullIntent: Intent = startedIntent!! + assertWithMessage("className") + .that(nonNullIntent.component?.className) + .isEqualTo(EvalService::class.java.name) + assertWithMessage("action") + .that(nonNullIntent.action) + .isEqualTo(EvalService.ACTION_START_SERVER) + assertWithMessage("port").that(nonNullIntent.getIntExtra("port", 0)).isEqualTo(9090) + } +} diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalServiceTest.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalServiceTest.kt new file mode 100644 index 000000000..4c57085be --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/EvalServiceTest.kt @@ -0,0 +1,137 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.content.Intent +import android.os.Build +import androidx.test.core.app.ApplicationProvider +import com.google.common.truth.Truth.assertThat +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.Robolectric +import org.robolectric.RobolectricTestRunner +import org.robolectric.Shadows.shadowOf +import org.robolectric.annotation.Config + +/** Unit tests for [EvalService] to verify foreground service lifecycle actions. */ +@RunWith(RobolectricTestRunner::class) +class EvalServiceTest { + + @Test + fun onStartCommand_startServer_startsForeground() { + val intent = + Intent(ApplicationProvider.getApplicationContext(), EvalService::class.java).apply { + action = EvalService.ACTION_START_SERVER + putExtra("port", 8083) + putExtra("model_path", "test_model.tflite") + putExtra("support_image", true) + putExtra("support_audio", true) + putExtra("accelerator", "CPU") + } + + val controller = Robolectric.buildService(EvalService::class.java, intent) + controller.create().startCommand(0, 1) + + val service = controller.get() + assertThat(service).isNotNull() + val shadowService = shadowOf(service) + assertThat(shadowService.lastForegroundNotificationId).isEqualTo(1) + assertThat(shadowService.lastForegroundNotification).isNotNull() + + // Test onBind for coverage + assertThat(service.onBind(intent)).isNull() + + // Test double start to cover "Server already running" branch + controller.startCommand(0, 2) + + // Test onDestroy + controller.destroy() + } + + @org.robolectric.annotation.Config(sdk = [33]) + @Test + fun onStartCommand_oldSdk_startsForeground() { + val intent = + Intent(ApplicationProvider.getApplicationContext(), EvalService::class.java).apply { + action = null + putExtra("port", 8083) + } + val controller = Robolectric.buildService(EvalService::class.java, intent) + controller.create().startCommand(0, 1) + + val service = controller.get() + val shadowService = shadowOf(service) + assertThat(shadowService.lastForegroundNotificationId).isEqualTo(1) + controller.destroy() + } + + @Test + fun onStartCommand_stopServer_stopsService() { + val intent = + Intent(ApplicationProvider.getApplicationContext(), EvalService::class.java).apply { + action = EvalService.ACTION_STOP_SERVER + } + + val controller = Robolectric.buildService(EvalService::class.java, intent) + controller.create().startCommand(0, 1) + + val service = controller.get() + assertThat(service).isNotNull() + assertThat(shadowOf(service).isStoppedBySelf).isTrue() + } + + @Test + fun onStartCommand_nullAction_startsForeground() { + val intent = Intent(ApplicationProvider.getApplicationContext(), EvalService::class.java) + intent.putExtra("model_path", "/fake/path") + intent.putExtra("support_audio", true) + intent.putExtra("support_image", true) + intent.putExtra("accelerator", "GPU") + val controller = Robolectric.buildService(EvalService::class.java, intent) + controller.create().startCommand(0, 1) + + val service = controller.get() + assertThat(shadowOf(service).lastForegroundNotificationId).isEqualTo(1) + } + + @Test + @Config(sdk = [Build.VERSION_CODES.TIRAMISU]) + fun onStartCommand_olderSdk_startsForeground() { + val intent = + Intent(ApplicationProvider.getApplicationContext(), EvalService::class.java).apply { + action = EvalService.ACTION_START_SERVER + } + val controller = Robolectric.buildService(EvalService::class.java, intent) + controller.create().startCommand(0, 1) + + val service = controller.get() + assertThat(shadowOf(service).lastForegroundNotificationId).isEqualTo(1) + } + + @Test + fun onDestroy_stopsServer() { + val intent = + Intent(ApplicationProvider.getApplicationContext(), EvalService::class.java).apply { + action = EvalService.ACTION_START_SERVER + } + val controller = Robolectric.buildService(EvalService::class.java, intent) + controller.create().startCommand(0, 1) + + // Trigger onDestroy + controller.destroy() + // It should not throw and server=null should be covered. + } +} diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/MiniHttpServerTest.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/MiniHttpServerTest.kt new file mode 100644 index 000000000..f1437409d --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/MiniHttpServerTest.kt @@ -0,0 +1,105 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import com.google.common.truth.Truth.assertThat +import java.io.BufferedReader +import java.io.InputStreamReader +import java.net.HttpURLConnection +import java.net.URL +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner + +/** + * Unit tests for [MiniHttpServer] to verify socket handling, request parsing, and routing logic. + */ +@RunWith(RobolectricTestRunner::class) +class MiniHttpServerTest { + + private lateinit var server: MiniHttpServer + private val port = 8081 + + @Before + fun setUp() { + server = + MiniHttpServer(port) { request -> + if (request.path == "/test") { + MiniHttpServer.Response(MiniHttpServer.Status.OK, "text/plain", "Test OK") + } else { + MiniHttpServer.Response(MiniHttpServer.Status.NOT_FOUND, "text/plain", "Not Found") + } + } + server.start() + } + + @After + fun tearDown() { + server.stop() + } + + @Test + fun handleRequest_validRoute_returnsOk() { + val url = URL("http://localhost:$port/test") + val connection = url.openConnection() as HttpURLConnection + connection.requestMethod = "GET" + + assertThat(connection.responseCode).isEqualTo(200) + + val response = BufferedReader(InputStreamReader(connection.inputStream)).readText() + assertThat(response).isEqualTo("Test OK") + } + + @Test + fun handleRequest_invalidRoute_returnsNotFound() { + val url = URL("http://localhost:$port/unknown") + val connection = url.openConnection() as HttpURLConnection + connection.requestMethod = "GET" + + assertThat(connection.responseCode).isEqualTo(404) + } + + @Test + fun handleRequest_emptyRequest_returnsEarly() { + val socket = java.net.Socket("localhost", port) + socket.getOutputStream().close() // Immediately close + socket.close() // Should return early without throwing + } + + @Test + fun handleRequest_malformedRequest_returnsEarly() { + val socket = java.net.Socket("localhost", port) + val out = socket.getOutputStream() + out.write("GET\r\n\r\n".toByteArray()) + out.flush() + socket.close() + } + + @Test + fun handleRequest_headersWithEmptyLinesAndContentLength_doesNotCrash() { + val socket = java.net.Socket("localhost", port) + val out = socket.getOutputStream() + val req = "POST /test HTTP/1.1\r\nContent-Length: 4\r\n\r\n\r\nbody" + out.write(req.toByteArray()) + out.flush() + val response = + java.io.BufferedReader(java.io.InputStreamReader(socket.getInputStream())).readText() + assertThat(response).contains("200 OK") + socket.close() + } +} diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ModelManagerTest.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ModelManagerTest.kt new file mode 100644 index 000000000..d07dff5d0 --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/ModelManagerTest.kt @@ -0,0 +1,284 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import android.content.Context +import android.graphics.Bitmap +import androidx.test.core.app.ApplicationProvider +import com.google.ai.edge.gallery.data.ConfigKeys +import com.google.ai.edge.gallery.data.Model +import com.google.ai.edge.gallery.data.RuntimeType +import com.google.ai.edge.gallery.runtime.CleanUpListener +import com.google.ai.edge.gallery.runtime.LlmModelHelper +import com.google.ai.edge.gallery.runtime.ResultListener +import com.google.ai.edge.litertlm.Contents +import com.google.ai.edge.litertlm.Message +import com.google.ai.edge.litertlm.ToolProvider +import com.google.async.coroutines.testing.doBlocking +import com.google.common.truth.Truth.assertThat +import java.util.concurrent.atomic.AtomicInteger +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Deferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.async +import kotlinx.coroutines.awaitAll +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner + +@RunWith(RobolectricTestRunner::class) +class ModelManagerTest { + + private lateinit var context: Context + private lateinit var modelManager: ModelManager + + @Before + fun setUp() { + context = ApplicationProvider.getApplicationContext() + modelManager = ModelManager(context) + } + + @Test + fun resolveModel_aicoreName_returnsAicoreModel() { + val model = + modelManager.resolveModel( + "some-aicore-model", + supportImage = true, + supportAudio = false, + accelerator = "GPU", + ) + assertThat(model).isNotNull() + assertThat(model!!.runtimeType).isEqualTo(RuntimeType.AICORE) + assertThat(model.llmSupportImage).isTrue() + assertThat(model.llmSupportAudio).isFalse() + assertThat(model.configValues[ConfigKeys.ACCELERATOR.label]).isEqualTo("GPU") + } + + @Test + fun resolveModel_localPath_returnsLitertModelWithPath() { + val model = + modelManager.resolveModel( + "/path/to/model.bin", + supportImage = true, + supportAudio = true, + accelerator = "GPU", + ) + assertThat(model).isNotNull() + assertThat(model!!.runtimeType).isEqualTo(RuntimeType.LITERT_LM) + assertThat(model.localModelFilePathOverride).isEqualTo("/path/to/model.bin") + assertThat(model.llmSupportImage).isTrue() + assertThat(model.llmSupportAudio).isTrue() + assertThat(model.configValues[ConfigKeys.ACCELERATOR.label]).isEqualTo("GPU") + } + + @Test + fun preInitModel_nullHelper_returnsEarly() { + val testManager = + object : ModelManager(context) { + override fun getHelperForModelName(modelName: String): LlmModelHelper? { + return null + } + } + // Should not crash, just returns early + testManager.preInitModel("test-pre-init", false, false, "CPU") + } + + @Test + fun getOrInitModel_success_cachesModel() = doBlocking { + val helper = FakeLlmModelHelper(initDelayMs = 10) + val modelName = "test-model" + + val model = + modelManager.getOrInitModel( + modelName, + helper, + supportImage = false, + supportAudio = false, + accelerator = "CPU", + ) + + assertThat(model).isNotNull() + assertThat(helper.initCallCount.get()).isEqualTo(1) + + // Second call should return cached model without calling initialize again + val model2 = + modelManager.getOrInitModel( + modelName, + helper, + supportImage = false, + supportAudio = false, + accelerator = "CPU", + ) + assertThat(model2).isSameInstanceAs(model) + assertThat(helper.initCallCount.get()).isEqualTo(1) + } + + @Test + fun getOrInitModel_concurrentRequests_initializesOnlyOnce() = doBlocking { + val helper = FakeLlmModelHelper(initDelayMs = 50) + val modelName = "concurrent-model" + + // Launch 3 concurrent requests + val deferreds = mutableListOf>() + val scope = CoroutineScope(Dispatchers.Default) + + for (i in 1..3) { + deferreds.add( + scope.async { + modelManager.getOrInitModel( + modelName, + helper, + supportImage = false, + supportAudio = false, + accelerator = "CPU", + ) + } + ) + } + + val results = deferreds.awaitAll() + + // Verify all got the same non-null model + assertThat(results[0]).isNotNull() + assertThat(results[1]).isSameInstanceAs(results[0]) + assertThat(results[2]).isSameInstanceAs(results[0]) + + // Verify initialize was only called once + assertThat(helper.initCallCount.get()).isEqualTo(1) + } + + @Test + fun getOrInitModel_failure_doesNotCache() = doBlocking { + val helper = FakeLlmModelHelper(shouldFail = true) + val modelName = "failing-model" + + val model = + modelManager.getOrInitModel( + modelName, + helper, + supportImage = false, + supportAudio = false, + accelerator = "CPU", + ) + assertThat(model).isNull() + assertThat(helper.initCallCount.get()).isEqualTo(1) + + // Second call should try to initialize again because first failed + val model2 = + modelManager.getOrInitModel( + modelName, + helper, + supportImage = false, + supportAudio = false, + accelerator = "CPU", + ) + assertThat(model2).isNull() + assertThat(helper.initCallCount.get()).isEqualTo(2) + } + + @Test + fun preInitModel_startsInitializationInBackground() = doBlocking { + val fakeHelper = FakeLlmModelHelper(initDelayMs = 10) + val testManager = + object : ModelManager(context) { + override fun getHelperForModelName(modelName: String): LlmModelHelper? { + return fakeHelper + } + } + + testManager.preInitModel( + "test-pre-init", + supportImage = false, + supportAudio = false, + accelerator = "CPU", + ) + + // Wait a bit for the background launch to complete + delay(50) + + assertThat(fakeHelper.initCallCount.get()).isEqualTo(1) + } + + // A simple fake implementation of LlmModelHelper for testing + class FakeLlmModelHelper( + private val initDelayMs: Long = 0, + private val shouldFail: Boolean = false, + ) : LlmModelHelper { + val initCallCount = AtomicInteger(0) + + override fun initialize( + context: Context, + model: Model, + taskId: String, + supportImage: Boolean, + supportAudio: Boolean, + onDone: (String) -> Unit, + systemInstruction: Contents?, + tools: List, + enableConversationConstrainedDecoding: Boolean, + coroutineScope: CoroutineScope?, + ) { + initCallCount.incrementAndGet() + val scope = coroutineScope ?: CoroutineScope(Dispatchers.Default) + scope.launch { + if (initDelayMs > 0) { + delay(initDelayMs) + } + if (shouldFail) { + onDone("Error: Simulated failure") + } else { + // Simulate the populated instance + model.instance = Any() // just a dummy object + onDone("Feature is available") + } + } + } + + override fun resetConversation( + model: Model, + supportImage: Boolean, + supportAudio: Boolean, + systemInstruction: Contents?, + tools: List, + enableConversationConstrainedDecoding: Boolean, + initialMessages: List, + ) {} + + override fun cleanUp(model: Model, onDone: () -> Unit) { + onDone() + } + + override fun runInference( + model: Model, + input: String, + resultListener: ResultListener, + cleanUpListener: CleanUpListener, + onError: (message: String) -> Unit, + images: List, + audioClips: List, + coroutineScope: CoroutineScope?, + extraContext: Map?, + ) { + // Simulate instantaneous inference completion to prevent tests from hanging + resultListener("mock response", true, null) + } + + override fun stopResponse(model: Model) {} + } +} diff --git a/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/PromptParserTest.kt b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/PromptParserTest.kt new file mode 100644 index 000000000..0a56d9bd4 --- /dev/null +++ b/Android/src/app/eval/javatests/com/google/ai/edge/gallery/eval/PromptParserTest.kt @@ -0,0 +1,158 @@ +/* + * Copyright 2024 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.ai.edge.gallery.eval + +import com.google.common.truth.Truth.assertThat +import org.json.JSONArray +import org.json.JSONObject +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.RobolectricTestRunner + +@RunWith(RobolectricTestRunner::class) +class PromptParserTest { + + // 1x1 pixel transparent GIF base64 + private val validBase64Image = "R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7" + private val dummyBase64Audio = "UklGRigAAABXQVZFZm10IBIAAAABAAEARKwAAIhYAQACABAAAABkYXRhAgAAAAAA" + + @Test + fun parseContent_textOnly_returnsText() { + val content = "Hello, world!" + val parsed = PromptParser.parseContent(content) + + assertThat(parsed.text).isEqualTo("Hello, world!") + assertThat(parsed.images).isEmpty() + assertThat(parsed.audioClips).isEmpty() + } + + @Test + fun parseContent_jsonArrayTextOnly_returnsText() { + val array = + JSONArray().apply { + put( + JSONObject().apply { + put("type", "text") + put("text", "Hello from JSON!") + } + ) + } + val parsed = PromptParser.parseContent(array) + + assertThat(parsed.text).isEqualTo("Hello from JSON!") + assertThat(parsed.images).isEmpty() + assertThat(parsed.audioClips).isEmpty() + } + + @Test + fun parseContent_multimodal_returnsParsedComponents() { + val array = + JSONArray().apply { + put( + JSONObject().apply { + put("type", "text") + put("text", "Describe this image and audio: ") + } + ) + put( + JSONObject().apply { + put("type", "image_url") + put( + "image_url", + JSONObject().apply { put("url", "data:image/gif;base64,$validBase64Image") }, + ) + } + ) + put( + JSONObject().apply { + put("type", "input_audio") + put("input_audio", JSONObject().apply { put("data", dummyBase64Audio) }) + } + ) + } + + val parsed = PromptParser.parseContent(array) + + assertThat(parsed.text).isEqualTo("Describe this image and audio: ") + assertThat(parsed.images).hasSize(1) + assertThat(parsed.images[0]).isNotNull() + assertThat(parsed.audioClips).hasSize(1) + assertThat(parsed.audioClips[0]).isNotEmpty() + } + + @Test + fun jsonToContents_textOnly_returnsTextContent() { + val content = "Simple text" + val lmContents = PromptParser.jsonToContents(content) + + // LmContents doesn't expose easy getters, but we can verify it doesn't crash + // and we can check its class type if needed. + assertThat(lmContents).isNotNull() + } + + @Test + fun convertToLmMessages_convertsCorrectly() { + val history = + listOf( + HistoryMessage("user", "Hello"), + HistoryMessage("assistant", "Hi there"), + HistoryMessage("system", "You are a helpful assistant"), + ) + + val lmMessages = PromptParser.convertToLmMessages(history) + + assertThat(lmMessages).hasSize(3) + // LmMessage has factory methods user(), model(), system(). + // We can't easily check contents without reflection, but we verify size and non-null. + assertThat(lmMessages[0]).isNotNull() + assertThat(lmMessages[1]).isNotNull() + assertThat(lmMessages[2]).isNotNull() + } + + @Test + fun convertToLmMessages_malformedBase64_doesNotCrash() { + val malformedHistory = + listOf( + HistoryMessage( + "user", + JSONArray() + .apply { + put( + JSONObject().apply { + put("type", "image_url") + put( + "image_url", + JSONObject().apply { put("url", "data:image/gif;base64,invalid_base64") }, + ) + } + ) + put( + JSONObject().apply { + put("type", "input_audio") + put("input_audio", JSONObject().apply { put("data", "invalid_base64") }) + } + ) + } + .toString(), + ) + ) + + val lmMessages = PromptParser.convertToLmMessages(malformedHistory) + + assertThat(lmMessages).hasSize(1) + assertThat(lmMessages[0]).isNotNull() + } +}