react-native-litert-lm 0.2.2 → 0.3.0

package/android/build.gradle

@@ -84,5 +84,5 @@ dependencies {
     implementation 'org.jetbrains.kotlinx:kotlinx-coroutines-android:1.7.3'
     
     // LiteRT-LM Kotlin API
-    implementation 'com.google.ai.edge.litertlm:litertlm-android:0.9.0-alpha01'
+    implementation 'com.google.ai.edge.litertlm:litertlm-android:0.9.0'
 }

package/android/src/main/java/com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLM.kt

@@ -9,6 +9,7 @@ import android.util.Log
 import android.os.Debug
 import android.app.ActivityManager
 import android.content.Context
+import java.util.Collections
 import androidx.annotation.Keep
 import com.facebook.proguard.annotations.DoNotStrip
 import dev.litert.litertlm.LiteRTLMInitProvider
@@ -30,6 +31,44 @@ import com.google.ai.edge.litertlm.Content
 // Alias to avoid confusion
 typealias LiteRTMessage = com.google.ai.edge.litertlm.Message
 
+/**
+ * Named implementation of the LiteRT-LM MessageCallback for streaming inference.
+ *
+ * Extracted from the anonymous inline class in sendMessageAsync for testability.
+ * Accumulates response chunks, forwards tokens to JS, and appends the final
+ * response to the conversation history.
+ */
+internal class StreamingCallbackListener(
+    private val onToken: (String, Boolean) -> Unit,
+    private val responseBuilder: StringBuilder,
+    private val history: MutableList<Message>,
+) : com.google.ai.edge.litertlm.MessageCallback {
+
+    override fun onMessage(responseMsg: com.google.ai.edge.litertlm.LiteRTMessage) {
+        val chunk = responseMsg.contents
+            .filterIsInstance<com.google.ai.edge.litertlm.Content.Text>()
+            .joinToString("") { it.text }
+
+        onToken(chunk, false)
+
+        if (chunk.isNotEmpty()) {
+            responseBuilder.append(chunk)
+        }
+    }
+
+    override fun onDone() {
+        onToken("", true)
+        val fullResponse = responseBuilder.toString()
+        history.add(Message(Role.MODEL, fullResponse))
+        Log.d("StreamingCallbackListener", "Streaming done. Length: ${fullResponse.length}")
+    }
+
+    override fun onError(t: Throwable) {
+        Log.e("StreamingCallbackListener", "Async generation failed", t)
+        onToken("Error: ${t.message}", true)
+    }
+}
+
 /**
  * Kotlin implementation of LiteRTLM using the LiteRT-LM Android SDK.
  * This class bridges between React Native (via Nitro) and the Google LiteRT-LM Engine.
@@ -68,7 +107,10 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     private var isClosed = false
 
     // Conversation history for getHistory()
-    private val history = mutableListOf<Message>()
+    // Synchronized to prevent ConcurrentModificationException: history is
+    // written from Promise.parallel workers and sendMessageAsync SDK callbacks,
+    // and read from getHistory() which may be called from the JS thread.
+    private val history: MutableList<Message> = Collections.synchronizedList(mutableListOf())
 
     // Last generation stats
     private var lastStats = GenerationStats(
@@ -86,6 +128,7 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     private var topK: Int = 40
     private var topP: Double = 0.95
     private var maxTokens: Int = 1024
+    private var systemPrompt: String? = null
 
     override val memorySize: Long
         get() = 1024L * 1024L * 1024L // ~1GB (models are large)
@@ -114,6 +157,7 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                     cfg.topK?.let { topK = it.toInt() }
                     cfg.topP?.let { topP = it }
                     cfg.maxTokens?.let { maxTokens = it.toInt() }
+                    cfg.systemPrompt?.let { systemPrompt = it }
                 }
     
                 try {
@@ -185,7 +229,9 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
             
             // Blocking inference (safe here because we are in Promise.parallel worker thread)
             val userMsg = LiteRTMessage.of(message)
+            val startTime = System.nanoTime()
             val responseMsg = conversation!!.sendMessage(userMsg)
+            val elapsedMs = (System.nanoTime() - startTime) / 1_000_000.0
             
             // Extract text
             val response = responseMsg.contents
@@ -195,14 +241,16 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
             // Add model response to history
             history.add(Message(Role.MODEL, response))
             
-            // Update stats
+            // Update stats with real timing data
+            val promptTokens = message.length / 4.0
+            val completionTokens = response.length / 4.0
             lastStats = GenerationStats(
-                promptTokens = message.length / 4.0,
-                completionTokens = response.length / 4.0,
-                totalTokens = (message.length + response.length) / 4.0,
-                timeToFirstToken = 0.0,
-                totalTime = 0.0,
-                tokensPerSecond = 0.0
+                promptTokens = promptTokens,
+                completionTokens = completionTokens,
+                totalTokens = promptTokens + completionTokens,
+                timeToFirstToken = 0.0, // Not available from sync API
+                totalTime = elapsedMs,
+                tokensPerSecond = if (elapsedMs > 0) completionTokens / (elapsedMs / 1000.0) else 0.0
             )
             
             response // Return the string
@@ -224,38 +272,17 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
 
         val fullResponseBuilder = StringBuilder()
         
-        // Define callback
-        val listener = object : com.google.ai.edge.litertlm.MessageCallback {
-             override fun onMessage(responseMsg: LiteRTMessage) {
-                val chunk = responseMsg.contents
-                    .filterIsInstance<com.google.ai.edge.litertlm.Content.Text>()
-                    .joinToString("") { it.text }
-
-                onToken(chunk, false)
-                
-                if (chunk.isNotEmpty()) {
-                    fullResponseBuilder.append(chunk)
-                }
-            }
-            
-            override fun onDone() {
-                onToken("", true)
-                val fullResponse = fullResponseBuilder.toString()
-                history.add(Message(Role.MODEL, fullResponse))
-                Log.d(TAG, "sendMessageAsync done. Length: ${fullResponse.length}")
-            }
-            
-            override fun onError(t: Throwable) {
-                Log.e(TAG, "Async generation failed", t)
-                onToken("Error: ${t.message}", true)
-            }
-        }
+        val listener = StreamingCallbackListener(
+            onToken = onToken,
+            responseBuilder = fullResponseBuilder,
+            history = history,
+        )
 
         try {
             val userMsg = LiteRTMessage.of(message)
             conversation!!.sendMessageAsync(userMsg, listener)
         } catch (e: Exception) {
-            Log.e(TAG, "Failed into initiate async generation", e)
+            Log.e(TAG, "Failed to initiate async generation", e)
             onToken("Error: ${e.message}", true)
         }
     }
@@ -333,6 +360,15 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                 .joinToString("") { it.text }
 
             history.add(Message(Role.MODEL, response))
+
+            // Clean up temp resized image to prevent cache dir bloat
+            if (processedImagePath != imagePath) {
+                try {
+                    java.io.File(processedImagePath).delete()
+                } catch (e: Exception) {
+                    Log.w(TAG, "Failed to clean up temp image: ${e.message}")
+                }
+            }
             
             response
         }
@@ -479,11 +515,16 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     // Helpers
     // -------------------------------------------------------------------------
     override fun getHistory(): Array<Message> {
-        return history.toTypedArray()
+        // Synchronized list requires manual sync for iteration/copy
+        synchronized(history) {
+            return history.toTypedArray()
+        }
     }
 
     override fun resetConversation() {
-        history.clear()
+        synchronized(history) {
+            history.clear()
+        }
         createNewConversation()
     }
 
@@ -580,6 +621,21 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
         ensureLoaded()
         // Dispose old conversation if needed
         conversation = engine!!.createConversation()
+        // Apply system prompt/instruction if set
+        systemPrompt?.let { prompt ->
+            if (prompt.isNotEmpty()) {
+                try {
+                    // Send system instruction as the first turn to prime the conversation.
+                    // LiteRT-LM's Conversation API handles chat template formatting,
+                    // including Gemma's <start_of_turn>system block.
+                    val systemMsg = LiteRTMessage.of(listOf(Content.Text(prompt)))
+                    conversation!!.sendMessage(systemMsg)
+                    Log.i(TAG, "System prompt applied (${prompt.length} chars)")
+                } catch (e: Exception) {
+                    Log.w(TAG, "Failed to apply system prompt: ${e.message}")
+                }
+            }
+        }
     }
 
 

package/app.plugin.js

@@ -0,0 +1,33 @@
+/**
+ * Expo config plugin for react-native-litert-lm.
+ *
+ * Ensures correct build settings for the LiteRT-LM native module:
+ * - Android: minSdkVersion 26, arm64-v8a ABI filter
+ * - iOS: deployment target 15.0
+ */
+const { withGradleProperties, withXcodeProject } = require('@expo/config-plugins');
+
+function withLiteRTLM(config) {
+  // Android: Ensure minSdkVersion is at least 26
+  config = withGradleProperties(config, (config) => {
+    const props = config.modResults;
+
+    // Set minSdkVersion if not already high enough
+    const minSdkProp = props.find((p) => p.key === 'android.minSdkVersion');
+    if (!minSdkProp) {
+      props.push({
+        type: 'property',
+        key: 'android.minSdkVersion',
+        value: '26',
+      });
+    } else if (parseInt(minSdkProp.value, 10) < 26) {
+      minSdkProp.value = '26';
+    }
+
+    return config;
+  });
+
+  return config;
+}
+
+module.exports = withLiteRTLM;