Fix API correctness issues found during static audit

- GeminiNanoModel: Rewrite to use actual ML Kit Prompt API (Generation.getClient(), checkStatus(), download(), warmup()) - MediaPipeModel: Fix streaming to use sync fallback since MediaPipe requires result listener set at build time - ApiModels: Add @SerializedName("object") for OpenAI compat, add "created" timestamp to ChatResponse/StreamChunk - settings.gradle.kts: Fix dependencyResolutionManagement typo - MainActivity: Use GradientDrawable.setColor() to preserve oval shape on status dot - Add Gradle wrapper scripts (gradlew, gradlew.bat, jar) https://claude.ai/code/session_01GvqMLSMmfMR8uz66BFVXX2
2026-02-28 14:27:46 +00:00
parent 12de200450
commit 43fb67d80a
8 changed files with 304 additions and 78 deletions
@@ -2,17 +2,16 @@ package com.pixel10.ai.inference

 import android.content.Context
 import android.util.Log
+import com.google.mlkit.genai.prompt.DownloadStatus
+import com.google.mlkit.genai.prompt.FeatureStatus
+import com.google.mlkit.genai.prompt.Generation
 import com.google.mlkit.genai.prompt.GenerativeModel
-import com.google.mlkit.genai.prompt.type.Content
 import com.google.mlkit.genai.prompt.type.TextPart
-import com.google.mlkit.genai.prompt.type.content
-import com.google.mlkit.genai.prompt.type.generationConfig
+import com.google.mlkit.genai.prompt.type.generateContentRequest
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.fold
-import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.withContext
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException

 /**
 * Gemini Nano backend via ML Kit Prompt API.
@@ -22,10 +21,9 @@ import kotlin.coroutines.resumeWithException
 * no manual download or file management needed.
 *
 * Key advantages:
- *  - Hardware-accelerated on Tensor G5 TPU (2.6x faster than G4)
- *  - 32,000 token context window on Pixel 10
- *  - ~3 GB model always resident in RAM for instant inference
+ *  - Hardware-accelerated on Tensor G5 TPU
 *  - Fully offline, private — data never leaves the device
+ *  - System-managed model, no manual downloads
 */
 class GeminiNanoModel private constructor(
    private val generativeModel: GenerativeModel
@@ -43,9 +41,12 @@ class GeminiNanoModel private constructor(
        temperature: Float
    ): String = withContext(Dispatchers.Default) {
        try {
-            val request = content { text(prompt) }
+            val request = generateContentRequest(TextPart(prompt)) {
+                this.temperature = temperature
+                this.topK = 40
+            }
            val response = generativeModel.generateContent(request)
-            response.text ?: ""
+            response.candidates.firstOrNull()?.text ?: ""
        } catch (e: Exception) {
            Log.e(TAG, "Gemini Nano generation error", e)
            throw OnDeviceModel.InferenceException("Gemini Nano generation failed: ${e.message}", e)
@@ -57,12 +58,11 @@ class GeminiNanoModel private constructor(
        onToken: (String) -> Unit
    ): String = withContext(Dispatchers.Default) {
        try {
-            val request = content { text(prompt) }
-            generativeModel.generateContentStream(request)
-                .fold("") { acc, response ->
-                    val chunk = response.text ?: ""
-                    if (chunk.isNotEmpty()) onToken(chunk)
-                    acc + chunk
+            generativeModel.generateContentStream(prompt)
+                .fold("") { acc, chunk ->
+                    val text = chunk.candidates.firstOrNull()?.text ?: ""
+                    if (text.isNotEmpty()) onToken(text)
+                    acc + text
                }
        } catch (e: Exception) {
            Log.e(TAG, "Gemini Nano streaming error", e)
@@ -72,53 +72,60 @@ class GeminiNanoModel private constructor(

    override fun close() {
        isReady = false
-        generativeModel.close()
+        // GenerativeModel from Generation.getClient() is system-managed
    }

    companion object {
        private const val TAG = "GeminiNanoModel"

        suspend fun create(context: Context): GeminiNanoModel = withContext(Dispatchers.IO) {
-            // Check if Gemini Nano is available on this device
-            val model = GenerativeModel.newBuilder()
-                .setContext(context)
-                .build()
+            val model = Generation.getClient()

-            // Verify feature is available — will throw if not supported
-            suspendCancellableCoroutine { continuation ->
-                model.isAvailable()
-                    .addOnSuccessListener { available ->
-                        if (available) {
-                            continuation.resume(Unit)
-                        } else {
-                            continuation.resumeWithException(
-                                OnDeviceModel.InferenceException(
-                                    "Gemini Nano is not available on this device"
-                                )
-                            )
+            // Check if Gemini Nano is available on this device
+            val status = model.checkStatus()
+            when (status) {
+                FeatureStatus.UNAVAILABLE -> {
+                    throw OnDeviceModel.InferenceException(
+                        "Gemini Nano is not available on this device"
+                    )
+                }
+                FeatureStatus.DOWNLOADABLE -> {
+                    Log.i(TAG, "Downloading Gemini Nano model...")
+                    model.download().collect { downloadStatus ->
+                        when (downloadStatus) {
+                            is DownloadStatus.DownloadStarted ->
+                                Log.i(TAG, "Model download started")
+                            is DownloadStatus.DownloadProgress ->
+                                Log.i(TAG, "Download in progress...")
+                            DownloadStatus.DownloadCompleted ->
+                                Log.i(TAG, "Model download completed")
+                            is DownloadStatus.DownloadFailed ->
+                                throw OnDeviceModel.InferenceException("Model download failed")
                        }
                    }
-                    .addOnFailureListener { e ->
-                        continuation.resumeWithException(
-                            OnDeviceModel.InferenceException(
-                                "Failed to check Gemini Nano availability: ${e.message}", e
-                            )
-                        )
+                }
+                FeatureStatus.DOWNLOADING -> {
+                    Log.i(TAG, "Model already downloading, waiting...")
+                    model.download().collect { downloadStatus ->
+                        if (downloadStatus == DownloadStatus.DownloadCompleted) {
+                            Log.i(TAG, "Download completed")
+                        }
                    }
+                }
+                FeatureStatus.AVAILABLE -> {
+                    Log.i(TAG, "Gemini Nano is available")
+                }
            }

-            // Trigger model download if needed
-            suspendCancellableCoroutine { continuation ->
-                model.downloadModel()
-                    .addOnSuccessListener { continuation.resume(Unit) }
-                    .addOnFailureListener { e ->
-                        Log.w(TAG, "Model download issue (may already be available): ${e.message}")
-                        // Don't fail — model might already be cached
-                        continuation.resume(Unit)
-                    }
+            // Warm up for lower first-inference latency
+            try {
+                model.warmup()
+                Log.i(TAG, "Model warmup complete")
+            } catch (e: Exception) {
+                Log.w(TAG, "Warmup failed (non-fatal): ${e.message}")
            }

-            Log.i(TAG, "Gemini Nano model ready via ML Kit Prompt API")
+            Log.i(TAG, "Gemini Nano ready via ML Kit Prompt API")
            GeminiNanoModel(model)
        }
    }
@@ -4,11 +4,8 @@ import android.content.Context
 import android.util.Log
 import com.google.mediapipe.tasks.genai.llminference.LlmInference
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.withContext
 import java.io.File
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException

 /**
 * MediaPipe LLM Inference backend for custom open-weight models.
@@ -52,23 +49,18 @@ class MediaPipeModel private constructor(
        prompt: String,
        onToken: (String) -> Unit
    ): String = withContext(Dispatchers.Default) {
-        suspendCancellableCoroutine { continuation ->
-            val fullResponse = StringBuilder()
-            try {
-                llmInference.generateResponseAsync(prompt).addResultListener { partialResult, done ->
-                    val chunk = partialResult ?: ""
-                    fullResponse.append(chunk)
-                    onToken(chunk)
-                    if (done) {
-                        continuation.resume(fullResponse.toString())
-                    }
-                }
-            } catch (e: Exception) {
-                Log.e(TAG, "Streaming error", e)
-                continuation.resumeWithException(
-                    OnDeviceModel.InferenceException("Streaming failed: ${e.message}", e)
-                )
-            }
+        // MediaPipe's streaming API (generateResponseAsync) requires the result
+        // listener to be set at LlmInference build time via setResultListener().
+        // Since our architecture needs a dynamic callback per request, we use
+        // synchronous generation and emit the result as a single chunk.
+        // For true token-by-token streaming, the Gemini Nano backend is preferred.
+        try {
+            val result = llmInference.generateResponse(prompt)
+            onToken(result)
+            result
+        } catch (e: Exception) {
+            Log.e(TAG, "MediaPipe generation error", e)
+            throw OnDeviceModel.InferenceException("Generation failed: ${e.message}", e)
        }
    }

@@ -120,7 +112,6 @@ class MediaPipeModel private constructor(
        }

        private fun findModelPath(context: Context): String? {
-            // Search standard locations
            val searchDirs = listOfNotNull(
                context.filesDir,
                File(context.filesDir, "models"),
@@ -137,7 +128,6 @@ class MediaPipeModel private constructor(
                        return file.absolutePath
                    }
                }
-                // Also check for any .task or .bin file
                dir.listFiles()?.firstOrNull {
                    it.extension in listOf("task", "bin", "tflite")
                }?.let { return it.absolutePath }
@@ -1,5 +1,7 @@
 package com.pixel10.ai.server

+import com.google.gson.annotations.SerializedName
+
 /**
 * Request/response models for the AI API.
 * Follows an OpenAI-compatible schema for easy integration.
@@ -20,6 +22,9 @@ data class Message(

 data class ChatResponse(
    val id: String,
+    @SerializedName("object")
+    val objectType: String = "chat.completion",
+    val created: Long = System.currentTimeMillis() / 1000,
    val model: String = "pixel10-on-device",
    val choices: List<Choice>,
    val usage: Usage
@@ -39,6 +44,9 @@ data class Usage(

 data class StreamChunk(
    val id: String,
+    @SerializedName("object")
+    val objectType: String = "chat.completion.chunk",
+    val created: Long = System.currentTimeMillis() / 1000,
    val model: String = "pixel10-on-device",
    val choices: List<StreamChoice>
 )
@@ -56,12 +64,15 @@ data class Delta(

 data class ModelInfo(
    val id: String = "pixel10-on-device",
-    val object_type: String = "model",
+    @SerializedName("object")
+    val objectType: String = "model",
    val owned_by: String = "local-device",
    val description: String = "On-device AI model running on Pixel 10 Tensor G5 chip"
 )

 data class ModelList(
+    @SerializedName("object")
+    val objectType: String = "list",
    val data: List<ModelInfo> = listOf(ModelInfo())
 )

@@ -6,6 +6,7 @@ import android.content.Context
 import android.content.Intent
 import android.content.ServiceConnection
 import android.content.pm.PackageManager
+import android.graphics.drawable.GradientDrawable
 import android.net.wifi.WifiManager
 import android.os.Build
 import android.os.Bundle
@@ -139,7 +140,7 @@ class MainActivity : AppCompatActivity() {
        when (state) {
            ApiServerService.ServerState.STOPPED -> {
                binding.tvServerStatus.text = getString(R.string.server_status_stopped)
-                binding.viewStatusDot.setBackgroundColor(getColor(R.color.status_stopped))
+                (binding.viewStatusDot.background as? GradientDrawable)?.setColor(getColor(R.color.status_stopped))
                binding.tvServerUrl.text = "http://—"
                binding.tvModelStatus.text = "Model: not loaded"
                binding.btnToggle.text = getString(R.string.btn_start)
@@ -148,7 +149,7 @@ class MainActivity : AppCompatActivity() {
            }
            ApiServerService.ServerState.LOADING_MODEL -> {
                binding.tvServerStatus.text = getString(R.string.server_status_starting)
-                binding.viewStatusDot.setBackgroundColor(getColor(R.color.primary))
+                (binding.viewStatusDot.background as? GradientDrawable)?.setColor(getColor(R.color.primary))
                binding.tvModelStatus.text = getString(R.string.model_loading)
                binding.btnToggle.isEnabled = false
                binding.etPort.isEnabled = false
@@ -157,7 +158,7 @@ class MainActivity : AppCompatActivity() {
                val port = binding.etPort.text.toString()
                val ip = getLocalIpAddress()
                binding.tvServerStatus.text = getString(R.string.server_status_running)
-                binding.viewStatusDot.setBackgroundColor(getColor(R.color.status_running))
+                (binding.viewStatusDot.background as? GradientDrawable)?.setColor(getColor(R.color.status_running))
                binding.tvServerUrl.text = "http://$ip:$port"
                binding.tvModelStatus.text = getString(R.string.model_ready)
                binding.btnToggle.text = getString(R.string.btn_stop)
@@ -166,7 +167,7 @@ class MainActivity : AppCompatActivity() {
            }
            ApiServerService.ServerState.ERROR -> {
                binding.tvServerStatus.text = getString(R.string.server_status_error)
-                binding.viewStatusDot.setBackgroundColor(getColor(R.color.error))
+                (binding.viewStatusDot.background as? GradientDrawable)?.setColor(getColor(R.color.error))
                binding.tvModelStatus.text = getString(R.string.model_error)
                binding.btnToggle.text = getString(R.string.btn_start)
                binding.btnToggle.isEnabled = true